diff --git a/datasette_query_assistant/__init__.py b/datasette_query_assistant/__init__.py
index 0f4a241..80f4b23 100644
--- a/datasette_query_assistant/__init__.py
+++ b/datasette_query_assistant/__init__.py
@@ -31,6 +31,7 @@ def _related(conn):
             return get_related_tables(conn, table)
 
         tables = await db.execute_fn(_related)
+        tables.add(table)
         sql = SCHEMA_SQL_SPECIFIC.replace("PARAMS", ",".join("?" for _ in tables))
         return (await db.execute(sql, tuple(tables))).first()[0]
     else:
diff --git a/datasette_query_assistant/templates/query_assistant.html b/datasette_query_assistant/templates/query_assistant.html
index c61f471..e277e1f 100644
--- a/datasette_query_assistant/templates/query_assistant.html
+++ b/datasette_query_assistant/templates/query_assistant.html
@@ -20,13 +20,14 @@ <h1>Query assistant for {% if table %}{{ table }}{% else %}{{ database }}{% endi
     <input type="hidden" name="table" value="{{ table }}">
     <input type="submit" value="Submit">
   </p>
+</form>
 
 <details><summary>Schema that will be passed to the model</summary>
 <pre>{{ schema }}</pre>
 </details>
 
 <script>
-  document.querySelector('textarea[name=question]').focus();
+  document.querySelector('#id_question').focus();
 </script>
 
 {% endblock %}
diff --git a/tests/test_query_assistant.py b/tests/test_query_assistant.py
index 3c66c12..8918f89 100644
--- a/tests/test_query_assistant.py
+++ b/tests/test_query_assistant.py
@@ -1,16 +1,18 @@
 from datasette.app import Datasette
 from datasette_query_assistant import get_related_tables
+import pytest_asyncio
 import pytest
 import sqlite_utils
 
 
-@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-query-assistant" in installed_plugins
+@pytest_asyncio.fixture
+async def datasette():
+    ds = Datasette()
+    db = ds.add_memory_database("test")
+    await db.execute_write(
+        "create table if not exists foo (id integer primary key, name text)"
+    )
+    return ds
 
 
 def test_get_related_tables():
@@ -25,3 +27,25 @@ def test_get_related_tables():
     assert get_related_tables(db.conn, "foo.bar.baz") == set()
     assert get_related_tables(db.conn, "species") == {"species", "animals"}
     assert get_related_tables(db.conn, "animals") == {"species", "animals"}
+
+
+@pytest.mark.asyncio
+async def test_database_assistant_page(datasette):
+    response = await datasette.client.get("/test/-/assistant")
+    assert response.status_code == 200
+    assert "Query assistant for test" in response.text
+    assert (
+        "<pre>CREATE TABLE foo (id integer primary key, name text)</pre>"
+        in response.text
+    )
+
+
+@pytest.mark.asyncio
+async def test_table_assistant_page(datasette):
+    response = await datasette.client.get("/test/-/assistant?table=foo")
+    assert response.status_code == 200
+    assert "Query assistant for foo" in response.text
+    assert (
+        "<pre>CREATE TABLE foo (id integer primary key, name text)</pre>"
+        in response.text
+    )