220 lines
10 KiB
Python
220 lines
10 KiB
Python
"""Bounded archive queries and event-loop responsiveness with synthetic data."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import sqlite3
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
|
|
import httpx
|
|
from fastapi import FastAPI
|
|
|
|
from admin_backend import Database
|
|
from archive_api import register_archive_routes
|
|
from archive_store import ArchiveStore
|
|
|
|
|
|
class MeasuredDatabase(Database):
|
|
max_steps = 0
|
|
statements: list[str] | None = None
|
|
|
|
def connect(self):
|
|
db = super().connect()
|
|
if self.statements is not None:
|
|
db.set_trace_callback(self.statements.append)
|
|
if self.max_steps:
|
|
steps = 0
|
|
def progress():
|
|
nonlocal steps
|
|
steps += 1000
|
|
return steps > self.max_steps
|
|
db.set_progress_handler(progress, 1000)
|
|
return db
|
|
|
|
|
|
def seed_archive(database, count=120_000):
|
|
now = "2026-09-14T00:00:00.000+00:00"
|
|
with database.connect() as db:
|
|
for tenant in ("default", "other"):
|
|
db.execute("""INSERT INTO archive_source_account
|
|
(id,tenant_id,external_account_id,created_at,updated_at)
|
|
VALUES (?,?,?,?,?)""", (tenant, tenant, tenant, now, now))
|
|
db.executemany("""INSERT INTO archive_person
|
|
(id,tenant_id,display_name,created_at,updated_at) VALUES (?,?,?,?,?)""",
|
|
((f"p{i:06}", "other" if i%2 else "default", f"Person {i}", now, now)
|
|
for i in range(2000)))
|
|
db.executemany("""INSERT INTO archive_conversation
|
|
(id,tenant_id,source_account_id,external_id,name,last_message_at,created_at,updated_at)
|
|
VALUES (?,?,?,?,?,?,?,?)""",
|
|
((f"c{i:06}", "other" if i%2 else "default", "other" if i%2 else "default",
|
|
str(i), f"Conversation {i}", now, now, now) for i in range(2000)))
|
|
db.executemany("""INSERT INTO archive_message
|
|
(id,tenant_id,conversation_id,source_account_id,sender_person_id,
|
|
sent_at,dedup_key,created_at,updated_at,content,message_type)
|
|
VALUES (?,?,?,?,?,?,?,?,?,?,?)""",
|
|
((f"m{i:09}", "other" if i%2 else "default", f"c{i%2000:06}",
|
|
"other" if i%2 else "default", f"p{i%2000:06}", now, str(i), now, now,
|
|
f"Synthetic message {i}", "text") for i in range(count)))
|
|
db.executemany("""INSERT INTO archive_export_job
|
|
(id,tenant_id,status,formats_json,filters_json,cutoff_at,created_at)
|
|
VALUES (?,? ,'completed','["csv"]','{}',?,?)""",
|
|
((f"job{i:04}", "default", now, now) for i in range(100)))
|
|
|
|
|
|
class ArchivePerformanceTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.temp = tempfile.TemporaryDirectory()
|
|
cls.database = MeasuredDatabase(Path(cls.temp.name)/"archive.db")
|
|
cls.database.initialize("", seed_admin=False)
|
|
cls.store = ArchiveStore(cls.database)
|
|
cls.store.initialize()
|
|
seed_archive(cls.database)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
cls.temp.cleanup()
|
|
|
|
def tearDown(self):
|
|
self.database.max_steps = 0
|
|
self.database.statements = None
|
|
|
|
def test_page_queries_do_not_rescan_all_messages_per_item(self):
|
|
# VM work budget is deterministic across slow/fast test machines.
|
|
self.database.max_steps = 500_000
|
|
page = self.store.conversations()
|
|
self.assertEqual(len(page["items"]), 50)
|
|
self.assertTrue(page["has_more"])
|
|
self.assertTrue(all(x["message_count"] == 60 for x in page["items"]))
|
|
self.assertTrue(all(int(x["external_id"])%2 == 0 for x in page["items"]))
|
|
people = self.store.people(200)
|
|
self.assertEqual(len(people), 200)
|
|
self.assertTrue(all(x["message_count"] == 60 for x in people))
|
|
|
|
def test_cursor_ties_have_no_duplicates_or_missing_rows(self):
|
|
ids = []
|
|
cursor = ""
|
|
while True:
|
|
page = self.store.conversations(200, cursor)
|
|
ids.extend(item["id"] for item in page["items"])
|
|
if not page["has_more"]:
|
|
break
|
|
cursor = page["next_cursor"]
|
|
self.assertEqual(len(ids), 1000)
|
|
self.assertEqual(len(set(ids)), 1000)
|
|
self.assertEqual(ids, sorted(ids, reverse=True))
|
|
first = self.store.messages(ids[0], 30)
|
|
second = self.store.messages(ids[0], 30, first["next_cursor"])
|
|
self.assertEqual(len(first["items"])+len(second["items"]), 60)
|
|
self.assertFalse(second["has_more"])
|
|
self.assertFalse({x["id"] for x in first["items"]}&{x["id"] for x in second["items"]})
|
|
|
|
def test_stats_tenant_isolation_and_latest_timestamp(self):
|
|
result = self.store.stats()
|
|
self.assertEqual(result["messages"], 60_000)
|
|
self.assertEqual(result["conversations"], 1000)
|
|
self.assertEqual(result["people"], 1000)
|
|
self.assertEqual(self.store.stats("absent")["messages"], 0)
|
|
self.assertEqual(result["last_message_at"], "2026-09-14T00:00:00.000+00:00")
|
|
|
|
def test_export_list_uses_two_queries_for_one_hundred_jobs(self):
|
|
self.database.statements = []
|
|
self.assertEqual(len(self.store.export_jobs(100)), 100)
|
|
selects = [s for s in self.database.statements if s.lstrip().upper().startswith("SELECT")]
|
|
self.assertEqual(len(selects), 2)
|
|
self.assertEqual(self.store.export_jobs(100, "other"), [])
|
|
|
|
def test_export_order_uses_index_without_full_sort(self):
|
|
self.database.statements = []
|
|
iterator = self.store._message_rows({}, "9999", batch_size=20)
|
|
rows = [next(iterator) for _ in range(20)]
|
|
iterator.close()
|
|
query = next(s for s in self.database.statements if "AS conversation_name" in s)
|
|
with self.database.connect() as db:
|
|
plan = [row[3] for row in db.execute("EXPLAIN QUERY PLAN "+query)]
|
|
self.assertFalse(any("TEMP B-TREE" in detail for detail in plan), plan)
|
|
self.assertEqual(len(rows), 20)
|
|
|
|
|
|
class ArchiveConcurrencyTest(unittest.TestCase):
|
|
def test_slow_archive_io_does_not_block_unrelated_requests(self):
|
|
async def scenario():
|
|
principal = SimpleNamespace(id=1)
|
|
def require(*codes):
|
|
return lambda: principal
|
|
database = mock.Mock()
|
|
with mock.patch("archive_api.ArchiveStore") as factory:
|
|
store = factory.return_value
|
|
store.storage_config.return_value = {"enabled": False}
|
|
app = FastAPI()
|
|
register_archive_routes(app, database, lambda: principal, require,
|
|
lambda request: "127.0.0.1", lambda: {"tenant_id":"default"},
|
|
lambda principal, account_id: "default")
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
|
cases = [
|
|
("conversations", "GET", "/api/v2/archive/conversations", None, {"items":[]}),
|
|
("import_messages", "POST", "/api/v2/archive/imports/messages",
|
|
{"source_account":{},"messages":[{}]}, {}),
|
|
("test_storage", "POST", "/api/v2/archive/storage/test", None, {"ok":True}),
|
|
]
|
|
for method, verb, path, body, result in cases:
|
|
entered, release = threading.Event(), threading.Event()
|
|
def slow(*args, **kwargs):
|
|
entered.set()
|
|
release.wait(2)
|
|
return result
|
|
getattr(store, method).side_effect = slow
|
|
request = asyncio.create_task(client.request(verb, path, json=body))
|
|
try:
|
|
self.assertTrue(await asyncio.to_thread(entered.wait, 3))
|
|
response = await asyncio.wait_for(client.get("/api/v2/archive/storage"), 1)
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertFalse(request.done(), method+" blocked the event loop")
|
|
finally:
|
|
release.set()
|
|
await request
|
|
asyncio.run(scenario())
|
|
|
|
def test_queued_exports_do_not_exhaust_api_workers(self):
|
|
async def scenario():
|
|
principal = SimpleNamespace(id=1)
|
|
def require(*codes):
|
|
return lambda: principal
|
|
entered, release = threading.Event(), threading.Event()
|
|
calls = []
|
|
def export(job_id):
|
|
calls.append(job_id)
|
|
entered.set()
|
|
release.wait(3)
|
|
with mock.patch("archive_api.ArchiveStore") as factory:
|
|
store = factory.return_value
|
|
store.create_export_job.return_value = {"id":"job"}
|
|
store.run_export_job.side_effect = export
|
|
store.storage_config.return_value = {"enabled":False}
|
|
app = FastAPI()
|
|
register_archive_routes(app, mock.Mock(), lambda: principal, require,
|
|
lambda request:"127.0.0.1", lambda: {}, lambda p,a:"default")
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app),base_url="http://test") as client:
|
|
requests = [asyncio.create_task(client.post("/api/v2/archive/exports",json={"formats":["csv"]})) for _ in range(45)]
|
|
try:
|
|
self.assertTrue(await asyncio.to_thread(entered.wait,3))
|
|
response = await asyncio.wait_for(client.get("/api/v2/archive/storage"),1)
|
|
self.assertEqual(response.status_code,200)
|
|
self.assertEqual(len(calls),1)
|
|
finally:
|
|
release.set()
|
|
results = await asyncio.gather(*requests)
|
|
self.assertTrue(all(item.status_code==200 for item in results))
|
|
self.assertEqual(len(calls),45)
|
|
asyncio.run(scenario())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|