Test for enrichment, closes #12

pyproject.toml

@@ -25,7 +25,7 @@ CI = "https://github.com/datasette/datasette-enrichments-gpt/actions"
 enrichments_gpt = "datasette_enrichments_gpt"
 
 [project.optional-dependencies]
-test = ["pytest", "pytest-asyncio"]
+test = ["datasette-test", "pytest", "pytest-asyncio", "pytest-recording"]
 
 [tool.pytest.ini_options]
 asyncio_mode = "strict"

tests/cassettes/test_enrichments_gpt/test_enrichments_gpt.yaml

@@ -0,0 +1,88 @@
+interactions:
+- request:
+    body: '{"model": "gpt-3.5-turbo", "messages": [{"role": "system", "content": "haiku"},
+      {"role": "user", "content": "Hearst Castle"}], "max_tokens": 1000}'
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '146'
+      content-type:
+      - application/json
+      host:
+      - api.openai.com
+      user-agent:
+      - python-httpx/0.27.0
+    method: POST
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    body:
+      string: !!binary |
+        H4sIAAAAAAAAA1RQy07DQAy85yusPaeINi2QHuFSnkIIISpA1WbjJFs269XaFS/131HSF1x8mPGM
+        x/OTAChbqiko02gxbXCD/NIVZXvz0BRPd5jdP8/19fPZ/PzKhMfRt0o7BRVLNLJTHRlqg0Ox5De0
+        iagFO9fh6XA8muT5KO+Jlkp0nawOMsiOJgNZxYIGx8PRZKtsyBpkNYWXBADgp59dRl/ip5rCcbpD
+        WmTWNarpfglARXIdojSzZdFeVHogDXlB38e+1UtksQa0UHj1F9rZiqK3GoyzVcUpzFBHFrjQLA5f
+        /cyyUPyCj8ZywMhqa7ve53FUh0hFl92vnNvjlfWWm0VEzeS72ywUNvJ1AvDW/73694oKkdogC6F3
+        9NzXuLFTh6b/kOMtKSTaHfBsmGzzKf5iwXZRWV9jDNFuSqjCIivyyUmps5NCJevkFwAA//8DAPop
+        wVkOAgAA
+    headers:
+      CF-Cache-Status:
+      - DYNAMIC
+      CF-RAY:
+      - 87b288af2e441010-LAX
+      Cache-Control:
+      - no-cache, must-revalidate
+      Connection:
+      - keep-alive
+      Content-Encoding:
+      - gzip
+      Content-Type:
+      - application/json
+      Date:
+      - Sat, 27 Apr 2024 23:18:49 GMT
+      Server:
+      - cloudflare
+      Set-Cookie:
+      - __cf_bm=qceh_pEnFUMrqQHfZtyUgKX1ZKrviI9BcXhmW7NUq6M-1714259929-1.0.1.1-83qAkTYLPpO_mdk9i4UzmJWQ5Q.y.wk2k0atOAHtajrPjnpVB739BngotPFAyJYzPfpKnUzT1fx4zNQggeYX9Q;
+        path=/; expires=Sat, 27-Apr-24 23:48:49 GMT; domain=.api.openai.com; HttpOnly;
+        Secure; SameSite=None
+      - _cfuvid=XdUvkQvvoddubcI96vDyZ3.Sd6GOA5OA2lqxb3NZNT8-1714259929916-0.0.1.1-604800000;
+        path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None
+      Transfer-Encoding:
+      - chunked
+      access-control-allow-origin:
+      - '*'
+      alt-svc:
+      - h3=":443"; ma=86400
+      openai-model:
+      - gpt-3.5-turbo-0125
+      openai-organization:
+      - user-r3e61fpak04cbaokp5buoae4
+      openai-processing-ms:
+      - '356'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=15724800; includeSubDomains
+      x-ratelimit-limit-requests:
+      - '5000'
+      x-ratelimit-limit-tokens:
+      - '160000'
+      x-ratelimit-remaining-requests:
+      - '4999'
+      x-ratelimit-remaining-tokens:
+      - '158993'
+      x-ratelimit-reset-requests:
+      - 12ms
+      x-ratelimit-reset-tokens:
+      - 377ms
+      x-request-id:
+      - req_aded4d1fda80dc957399e53559e60fb3
+    status:
+      code: 200
+      message: OK
+version: 1

tests/conftest.py

@@ -0,0 +1,6 @@
+import pytest
+
+
+@pytest.fixture(scope="module")
+def vcr_config():
+    return {"filter_headers": ["authorization"]}

tests/test_enrichments_gpt.py

@@ -1,11 +1,58 @@
 from datasette.app import Datasette
+from datasette_test import actor_cookie
+from datasette_enrichments.utils import wait_for_job
+import os
 import pytest
 
 
 @pytest.mark.asyncio
-async def test_plugin_is_installed():
-    datasette = Datasette(memory=True)
-    response = await datasette.client.get("/-/plugins.json")
-    assert response.status_code == 200
-    installed_plugins = {p["name"] for p in response.json()}
-    assert "datasette-enrichments-gpt" in installed_plugins
+@pytest.mark.vcr(ignore_localhost=True)
+async def test_enrichments_gpt(monkeypatch):
+    if not os.environ.get("DATASETTE_SECRETS_OPENAI_API_KEY"):
+        monkeypatch.setenv("DATASETTE_SECRETS_OPENAI_API_KEY", "sk-xyz")
+    ds = Datasette()
+    db = ds.add_memory_database("test")
+    await db.execute_write("create table museums (id integer primary key, name text)")
+    await db.execute_write("insert into museums (name) values ('Hearst Castle')")
+    cookies = {"ds_actor": actor_cookie(ds, {"id": "root"})}
+
+    response1 = await ds.client.get("/-/enrich/test/museums/gpt", cookies=cookies)
+    assert "<h2>AI analysis with OpenAI GPT</h2>" in response1.text
+
+    # Now try and run it
+    csrftoken = response1.cookies["ds_csrftoken"]
+    cookies["ds_csrftoken"] = csrftoken
+
+    response2 = await ds.client.post(
+        "/-/enrich/test/museums/gpt",
+        cookies=cookies,
+        data={
+            "model": "gpt-3.5-turbo",
+            "prompt": "{{ name }}",
+            "system_prompt": "haiku",
+            "output_column": "haiku",
+            "csrftoken": csrftoken,
+        },
+    )
+    assert response2.status_code == 302
+    job_id = response2.headers["location"].split("=")[-1]
+    # Wait for it to finish
+    await wait_for_job(ds, job_id, "test", timeout=10)
+    # Should have no errors
+    errors = [
+        dict(row)
+        for row in await db.execute(
+            "select job_id, row_pks, error from _enrichment_errors"
+        )
+    ]
+    assert errors == []
+    # Check the result
+    response3 = await ds.client.get("/test/museums.json?_shape=array")
+    data = response3.json()
+    assert data == [
+        {
+            "id": 1,
+            "name": "Hearst Castle",
+            "haiku": "Majestic atop\nCalifornia cliffs, Hearst Castle\nHistory whispers",
+        }
+    ]