Skip to content

Commit

Permalink
continue improving tests, #22
Browse files Browse the repository at this point in the history
  • Loading branch information
jamesturk committed Mar 26, 2023
1 parent 7c5b520 commit d513cdb
Show file tree
Hide file tree
Showing 6 changed files with 88 additions and 38 deletions.
File renamed without changes.
File renamed without changes.
File renamed without changes.
39 changes: 1 addition & 38 deletions tests/test_apicall.py
Original file line number Diff line number Diff line change
@@ -1,45 +1,8 @@
import pytest
from unittest.mock import patch
from scrapeghost.apicall import OpenAiCall, RetryRule
from scrapeghost.errors import MaxCostExceeded
import openai


def _mock_response(**kwargs):
mr = openai.openai_object.OpenAIObject.construct_from(
dict(
object="text_completion",
model=kwargs.get("model"),
choices=[
{
"finish_reason": kwargs.get("finish_reason", "stop"),
"message": {
"content": kwargs.get("content", "Hello world"),
},
}
],
usage={
"prompt_tokens": kwargs.get("prompt_tokens", 1),
"completion_tokens": kwargs.get("completion_tokens", 1),
},
created=1629200000,
id="cmpl-xxxxxxxxxxxxxxxxxxxx",
model_version=kwargs.get("model_version", "ada"),
prompt="Hello world",
status="complete",
)
)

return mr


def _timeout(**kwargs):
raise openai.error.Timeout()


def patch_create():
p = patch("scrapeghost.apicall.openai.ChatCompletion.create")
return p
from testutils import _mock_response, _timeout, patch_create


def test_basic_call():
Expand Down
49 changes: 49 additions & 0 deletions tests/test_json_postprocessor.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
import pytest
from scrapeghost import SchemaScraper
from scrapeghost.postprocessors import JSONPostprocessor
from scrapeghost.responses import Response
from scrapeghost.errors import InvalidJSON, PostprocessingError
from testutils import patch_create, _mock_response


def test_json_already_processed():
# already processed
jpp = JSONPostprocessor()
r = Response(data={})
with pytest.raises(PostprocessingError):
jpp(r, scraper=SchemaScraper({"name": "string"}))


def test_json_no_nudge():
# single quotes, trailing commas
bad_json = "{'name': 'phil', }"
jpp = JSONPostprocessor(nudge=False)
r = Response(data=bad_json)
with pytest.raises(InvalidJSON):
jpp(r, scraper=SchemaScraper({"name": "string"}))
assert "False" in str(jpp)


def test_json_nudge():
# single quotes, trailing commas
bad_json = "{'name': 'phil', }"
jpp = JSONPostprocessor(nudge=True)
r = Response(data=bad_json)
with patch_create() as create:
create.side_effect = lambda **kwargs: _mock_response(content='{"name": "phil"}')
repaired = jpp(r, scraper=SchemaScraper({"name": "string"}))
assert len(repaired.api_responses) == 1
assert repaired.data == {"name": "phil"}


def test_nudge_fails():
# single quotes, trailing commas
bad_json = "{'name': 'phil', }"
jpp = JSONPostprocessor(nudge=True)
r = Response(data=bad_json)
with pytest.raises(InvalidJSON):
with patch_create() as create:
create.side_effect = lambda **kwargs: _mock_response(
content='{"name": "phil'
)
jpp(r, scraper=SchemaScraper({"name": "string"}))
38 changes: 38 additions & 0 deletions tests/testutils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
from unittest.mock import patch
import openai


def _mock_response(**kwargs):
mr = openai.openai_object.OpenAIObject.construct_from(
dict(
model=kwargs.get("model"),
choices=[
{
"finish_reason": kwargs.get("finish_reason", "stop"),
"message": {
"content": kwargs.get("content", "Hello world"),
},
}
],
usage={
"prompt_tokens": kwargs.get("prompt_tokens", 1),
"completion_tokens": kwargs.get("completion_tokens", 1),
},
created=1629200000,
id="cmpl-xxxxxxxxxxxxxxxxxxxx",
model_version=kwargs.get("model_version", "ada"),
prompt="Hello world",
status="complete",
)
)

return mr


def _timeout(**kwargs):
raise openai.error.Timeout()


def patch_create():
p = patch("scrapeghost.apicall.openai.ChatCompletion.create")
return p

0 comments on commit d513cdb

Please sign in to comment.