163 lines
5.3 KiB
Python
163 lines
5.3 KiB
Python
"""create_tables=True when several processes start against the same empty database.
|
|
|
|
Both processes see the tables missing and both CREATE them; one fails. The
|
|
sink retries create_all once, which skips what the other process made.
|
|
"""
|
|
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from sqlalchemy import create_engine, func, select
|
|
from sqlalchemy.exc import OperationalError
|
|
|
|
from fastmcp_feedback.instrumentation import (
|
|
DatabaseSink,
|
|
InstrumentationMiddleware,
|
|
build_metadata,
|
|
)
|
|
|
|
EVENTS = build_metadata().tables["ffb_events"]
|
|
WORKER = Path(__file__).parent.parent / "fixtures" / "create_tables_race_worker.py"
|
|
|
|
|
|
def _already_exists():
|
|
return OperationalError(
|
|
"CREATE TABLE ffb_tool_calls",
|
|
{},
|
|
Exception("table ffb_tool_calls already exists"),
|
|
)
|
|
|
|
|
|
def race(url: str, trigger: Path, workers: int = 2, events: int = 20) -> list[dict]:
|
|
"""Start the workers, release them together, return their JSON results."""
|
|
procs = [
|
|
subprocess.Popen(
|
|
[sys.executable, str(WORKER), url, str(trigger), str(events)],
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
for _ in range(workers)
|
|
]
|
|
try:
|
|
for p in procs:
|
|
assert p.stdout.readline().strip() == "ready"
|
|
trigger.touch()
|
|
results = []
|
|
for p in procs:
|
|
out, err = p.communicate(timeout=30)
|
|
assert p.returncode == 0, err
|
|
results.append(json.loads(out.strip().splitlines()[-1]))
|
|
return results
|
|
finally:
|
|
for p in procs:
|
|
if p.poll() is None:
|
|
p.kill()
|
|
p.wait()
|
|
|
|
|
|
@pytest.mark.slow
|
|
def test_two_processes_cold_start(tmp_path):
|
|
workers, events = 2, 20
|
|
for run in range(5):
|
|
db = tmp_path / f"shared{run}.db"
|
|
results = race(
|
|
f"sqlite+aiosqlite:///{db}", tmp_path / f"go{run}", workers, events
|
|
)
|
|
engine = create_engine(f"sqlite:///{db}")
|
|
try:
|
|
with engine.connect() as conn:
|
|
stored = conn.execute(
|
|
select(func.count()).select_from(EVENTS)
|
|
).scalar_one()
|
|
finally:
|
|
engine.dispose()
|
|
assert [r["sink_errors"] for r in results] == [0] * workers, f"run {run}"
|
|
assert stored == workers * events, f"run {run}"
|
|
|
|
|
|
@pytest.fixture
|
|
async def sink(tmp_path):
|
|
s = DatabaseSink(f"sqlite+aiosqlite:///{tmp_path}/calls.db", create_tables=True)
|
|
yield s
|
|
await s.engine.dispose()
|
|
|
|
|
|
class TestEnsureTablesRetry:
|
|
async def test_first_create_fails_retry_succeeds(self, sink, monkeypatch, caplog):
|
|
real = sink.metadata.create_all
|
|
calls = []
|
|
|
|
def create_all(bind, **kwargs):
|
|
calls.append(1)
|
|
if len(calls) == 1:
|
|
raise _already_exists()
|
|
return real(bind, **kwargs)
|
|
|
|
monkeypatch.setattr(sink.metadata, "create_all", create_all)
|
|
mw = InstrumentationMiddleware([sink])
|
|
caplog.set_level("DEBUG", logger="fastmcp_feedback.instrumentation.storage")
|
|
assert mw.record_event("a.b", key="k") is True
|
|
await mw.flush()
|
|
assert mw.dispatcher.sink_errors == 0
|
|
assert len(calls) == 2
|
|
assert sink._tables_ready
|
|
async with sink.engine.connect() as conn:
|
|
count = (
|
|
await conn.execute(select(func.count()).select_from(sink.events))
|
|
).scalar_one()
|
|
assert count == 1
|
|
retry_logs = [r for r in caplog.records if "retry succeeded" in r.getMessage()]
|
|
assert retry_logs and all(r.levelname == "INFO" for r in retry_logs)
|
|
assert not [
|
|
r
|
|
for r in caplog.records
|
|
if r.levelname == "WARNING" and r.name.endswith("storage")
|
|
]
|
|
|
|
async def test_persistent_failure_surfaces(self, sink, monkeypatch):
|
|
calls = []
|
|
|
|
def create_all(bind, **kwargs):
|
|
calls.append(1)
|
|
raise _already_exists()
|
|
|
|
monkeypatch.setattr(sink.metadata, "create_all", create_all)
|
|
mw = InstrumentationMiddleware([sink])
|
|
assert mw.record_event("a.b") is True
|
|
await mw.flush()
|
|
assert mw.dispatcher.sink_errors == 1
|
|
assert len(calls) == 2
|
|
assert not sink._tables_ready
|
|
# The next batch tries again rather than assuming the tables exist.
|
|
assert mw.record_event("a.c") is True
|
|
await mw.flush()
|
|
assert mw.dispatcher.sink_errors == 2
|
|
assert len(calls) == 4
|
|
await mw.aclose()
|
|
|
|
async def test_non_database_error_not_retried(self, sink, monkeypatch):
|
|
calls = []
|
|
|
|
def create_all(bind, **kwargs):
|
|
calls.append(1)
|
|
raise RuntimeError("bug")
|
|
|
|
monkeypatch.setattr(sink.metadata, "create_all", create_all)
|
|
with pytest.raises(RuntimeError):
|
|
await sink.ensure_tables()
|
|
assert len(calls) == 1
|
|
|
|
async def test_create_tables_false_never_creates(self, tmp_path, monkeypatch):
|
|
s = DatabaseSink(f"sqlite+aiosqlite:///{tmp_path}/none.db")
|
|
calls = []
|
|
monkeypatch.setattr(s.metadata, "create_all", lambda *a, **k: calls.append(1))
|
|
try:
|
|
await s._ready()
|
|
finally:
|
|
await s.engine.dispose()
|
|
assert calls == []
|