continue improving tests, #22

tests/test_apicall.py

@@ -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():

tests/test_json_postprocessor.py

@@ -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"}))

tests/testutils.py

@@ -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