Tests for table count endpoint, refs #2914

Forgot to include these tests in commit 1538832830

tests/test_table_count.py

@@ -0,0 +1,191 @@
+import pytest
+import pytest_asyncio
+
+from datasette import hookimpl
+from datasette.app import Datasette
+from datasette.database import QueryInterrupted
+from datasette.filters import FilterArguments
+from datasette.plugins import pm
+from datasette.utils import tilde_encode
+
+
+@pytest_asyncio.fixture
+async def count_ds(tmp_path):
+    ds = Datasette(
+        [tmp_path / "counts.db"],
+        config={"permissions": {"execute-sql": {"id": "sql-user"}}},
+    )
+    db = ds.get_database()
+    await db.execute_write_script("""
+        create table numbers (id integer primary key, name text);
+        insert into numbers values (1, 'one'), (2, 'two'), (3, 'three'),
+            (4, 'four'), (5, 'five'), (6, 'six'), (7, 'seven');
+        create view number_view as select id + 0 as number from numbers;
+        create table "a/b.c" (id integer);
+        insert into "a/b.c" values (1);
+        create virtual table numbers_fts using fts5(name, content="numbers");
+        insert into numbers_fts(numbers_fts) values ('rebuild');
+    """)
+    db.count_limit = 2
+    yield ds
+    ds.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "table,query,count",
+    [
+        ("numbers", "", 7),
+        ("numbers", "?id__gt=1", 6),
+        ("numbers", "?id__gt=1&id__lt=5", 3),
+        ("numbers", "?id__gt=1&id__gt=3", 4),
+        ("numbers", "?id__gt=100", 0),
+        ("numbers", "?name__contains=o", 3),
+        ("numbers", "?_search=two", 1),
+        ("numbers", "?_search_name=three", 1),
+        ("numbers", "?id__gt=1&_next=5&_size=1&_sort_desc=id&_nocount=1", 6),
+        # A computed column without affinity catches string/integer casting bugs.
+        ("number_view", "?number__gt=1", 6),
+        ("a/b.c", "", 1),
+    ],
+)
+async def test_count(count_ds, table, query, count):
+    response = await count_ds.client.post(
+        f"/counts/{tilde_encode(table)}/-/count{query}"
+    )
+    assert response.status_code == 200, response.text
+    assert response.json() == {"ok": True, "count": count}
+    assert response.headers["cache-control"] == "no-store"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["get", "head", "put", "delete"])
+async def test_count_post_only(count_ds, method):
+    response = await count_ds.client.request(method, "/counts/numbers/-/count")
+    assert response.status_code == 405
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "path,status",
+    [
+        ("/missing/numbers/-/count", 404),
+        ("/counts/missing/-/count", 404),
+        ("/counts/numbers/-/count?_fts_table=invalid", 400),
+        ("/counts/numbers/-/count?_where=id>1", 403),
+    ],
+)
+async def test_count_errors(count_ds, path, status):
+    response = await count_ds.client.post(
+        path, json={}, headers={"Accept": "application/json"}
+    )
+    assert response.status_code == status, response.text
+    assert response.json()["ok"] is False
+
+
+@pytest.mark.asyncio
+async def test_count_permissions(tmp_path):
+    ds = Datasette(
+        [tmp_path / "counts.db"], config={"permissions": {"view-table": False}}
+    )
+    await ds.get_database().execute_write("create table secret (id integer)")
+    try:
+        response = await ds.client.post(
+            "/counts/secret/-/count", json={}, headers={"Accept": "application/json"}
+        )
+        assert response.status_code == 403
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+async def test_count_plugin_filter(count_ds):
+    class Plugin:
+        @hookimpl
+        def filters_from_request(self, request):
+            if request.args.get("_custom"):
+                return FilterArguments(["id > :custom"], params={"custom": 4})
+
+    plugin = Plugin()
+    pm.register(plugin)
+    try:
+        response = await count_ds.client.post(
+            "/counts/numbers/-/count?_custom=1",
+            json={},
+            headers={"Accept": "application/json"},
+        )
+        assert response.json() == {"ok": True, "count": 3}
+    finally:
+        pm.unregister(plugin)
+
+
+@pytest.mark.asyncio
+async def test_count_timeout(count_ds, monkeypatch):
+    db = count_ds.get_database()
+    original = db.execute
+
+    async def execute(sql, *args, **kwargs):
+        if sql.startswith("select count(*)"):
+            raise QueryInterrupted(Exception("interrupted"), sql, {})
+        return await original(sql, *args, **kwargs)
+
+    monkeypatch.setattr(db, "execute", execute)
+    response = await count_ds.client.post(
+        "/counts/numbers/-/count", json={}, headers={"Accept": "application/json"}
+    )
+    assert response.status_code == 400
+    assert response.json()["errors"] == ["Count query timed out"]
+
+
+@pytest.mark.asyncio
+async def test_count_button_without_execute_sql(count_ds):
+    from bs4 import BeautifulSoup
+
+    response = await count_ds.client.get("/counts/numbers?id__gt=1")
+    soup = BeautifulSoup(response.text, "html.parser")
+    button = soup.select_one("button.count-all")
+    assert button["data-count-url"] == "/counts/numbers/-/count"
+    assert soup.select_one(".table-count").text == "2+ rows"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "query,status,count",
+    [
+        ("?_where=id>2&_where=id<6", 200, 3),
+        ("?_where=invalid_sql(", 400, None),
+    ],
+)
+async def test_count_where(count_ds, query, status, count):
+    response = await count_ds.client.post(
+        "/counts/numbers/-/count" + query,
+        json={},
+        actor={"id": "sql-user"},
+    )
+    assert response.status_code == status
+    if status == 200:
+        assert response.json() == {"ok": True, "count": count}
+    else:
+        assert response.json()["ok"] is False
+
+
+@pytest.mark.asyncio
+async def test_count_base_url(tmp_path):
+    from bs4 import BeautifulSoup
+
+    ds = Datasette([tmp_path / "counts.db"], settings={"base_url": "/prefix/"})
+    db = ds.get_database()
+    await db.execute_write_script(
+        "create table numbers (id integer); insert into numbers values (1), (2)"
+    )
+    db.count_limit = 1
+    try:
+        response = await ds.client.get("/counts/numbers")
+        button = BeautifulSoup(response.text, "html.parser").select_one(".count-all")
+        assert button["data-count-url"] == "/prefix/counts/numbers/-/count"
+        response = await ds.client.post(
+            ds.urls.path("/counts/numbers/-/count"), json={}
+        )
+        assert response.json() == {"ok": True, "count": 2}
+    finally:
+        ds.close()