llm_council/backend/tests/test_docs_api.py

import os
import shutil
import tempfile
import unittest
import importlib

import httpx


class TestDocumentsApi(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self._old_env = dict(os.environ)
        self.tmp = tempfile.mkdtemp(prefix="llm-council-docsapi-")
        os.environ["DOCS_DIR"] = self.tmp
        os.environ["MAX_DOC_BYTES"] = "1000000"

        # Reload config/documents so they see DOCS_DIR override
        import backend.config as config
        import backend.documents as documents
        import backend.main as main

        importlib.reload(config)
        importlib.reload(documents)
        self.main = importlib.reload(main)

        self.transport = httpx.ASGITransport(app=self.main.app)
        self.client = httpx.AsyncClient(transport=self.transport, base_url="http://test")

        # Create a conversation
        resp = await self.client.post("/api/conversations", json={})
        resp.raise_for_status()
        self.conversation_id = resp.json()["id"]

    async def asyncTearDown(self):
        await self.client.aclose()
        os.environ.clear()
        os.environ.update(self._old_env)
        shutil.rmtree(self.tmp, ignore_errors=True)

    async def test_upload_and_list_documents(self):
        # Upload
        files = {"file": ("notes.md", b"# Hi\n", "text/markdown")}
        resp = await self.client.post(
            f"/api/conversations/{self.conversation_id}/documents",
            files=files,
        )
        self.assertEqual(resp.status_code, 200, resp.text)
        meta = resp.json()
        self.assertIn("id", meta)
        self.assertEqual(meta["filename"], "notes.md")

        # List
        resp2 = await self.client.get(
            f"/api/conversations/{self.conversation_id}/documents"
        )
        self.assertEqual(resp2.status_code, 200, resp2.text)
        items = resp2.json()
        self.assertEqual(len(items), 1)
        self.assertEqual(items[0]["id"], meta["id"])

    async def test_upload_multiple_documents(self):
        files = [
            ("files", ("a.md", b"one", "text/markdown")),
            ("files", ("b.md", b"two", "text/markdown")),
        ]
        resp = await self.client.post(
            f"/api/conversations/{self.conversation_id}/documents",
            files=files,
        )
        self.assertEqual(resp.status_code, 200, resp.text)
        payload = resp.json()
        self.assertIn("uploaded", payload)
        self.assertEqual(len(payload["uploaded"]), 2)
        self.assertEqual({d["filename"] for d in payload["uploaded"]}, {"a.md", "b.md"})

    async def test_rejects_non_md(self):
        files = {"file": ("notes.txt", b"hello", "text/plain")}
        resp = await self.client.post(
            f"/api/conversations/{self.conversation_id}/documents",
            files=files,
        )
        self.assertEqual(resp.status_code, 400)

    async def test_get_and_delete_document(self):
        files = {"file": ("a.md", b"hello", "text/markdown")}
        up = await self.client.post(
            f"/api/conversations/{self.conversation_id}/documents",
            files=files,
        )
        self.assertEqual(up.status_code, 200)
        doc_id = up.json()["id"]

        get = await self.client.get(
            f"/api/conversations/{self.conversation_id}/documents/{doc_id}"
        )
        self.assertEqual(get.status_code, 200)
        self.assertEqual(get.json()["content"], "hello")

        dele = await self.client.delete(
            f"/api/conversations/{self.conversation_id}/documents/{doc_id}"
        )
        self.assertEqual(dele.status_code, 200)
        self.assertTrue(dele.json()["ok"])

        get2 = await self.client.get(
            f"/api/conversations/{self.conversation_id}/documents/{doc_id}"
        )
        self.assertEqual(get2.status_code, 404)