test(monitoring): cover trace observability

tests/unit_tests/persistence/test_monitoring_trace_migration.py

@@ -0,0 +1,87 @@
+"""Unit tests for the monitoring trace Alembic migration."""
+
+from __future__ import annotations
+
+from importlib import import_module
+
+
+class _FakeInspector:
+    def __init__(self, tables):
+        self._tables = tables
+
+    def get_table_names(self):
+        return list(self._tables)
+
+
+class _FakeOp:
+    def __init__(self):
+        self.created_tables = []
+        self.created_indexes = []
+        self.dropped_tables = []
+
+    def get_bind(self):
+        return object()
+
+    def create_table(self, table_name, *columns):
+        self.created_tables.append((table_name, columns))
+
+    def create_index(self, index_name, table_name, columns):
+        self.created_indexes.append((index_name, table_name, columns))
+
+    def drop_table(self, table_name):
+        self.dropped_tables.append(table_name)
+
+
+def _migration_module():
+    return import_module('langbot.pkg.persistence.alembic.versions.0006_monitoring_traces')
+
+
+def test_upgrade_creates_monitoring_trace_tables_and_indexes(monkeypatch):
+    migration = _migration_module()
+    fake_op = _FakeOp()
+
+    monkeypatch.setattr(migration, 'op', fake_op)
+    monkeypatch.setattr(migration.sa, 'inspect', lambda _conn: _FakeInspector(tables=set()))
+
+    migration.upgrade()
+
+    assert [table_name for table_name, _columns in fake_op.created_tables] == [
+        'monitoring_traces',
+        'monitoring_spans',
+    ]
+    assert ('ix_monitoring_traces_started_at', 'monitoring_traces', ['started_at']) in fake_op.created_indexes
+    assert ('ix_monitoring_spans_trace_id', 'monitoring_spans', ['trace_id']) in fake_op.created_indexes
+    assert ('ix_monitoring_spans_pipeline_id', 'monitoring_spans', ['pipeline_id']) in fake_op.created_indexes
+
+
+def test_upgrade_skips_existing_monitoring_trace_tables(monkeypatch):
+    migration = _migration_module()
+    fake_op = _FakeOp()
+
+    monkeypatch.setattr(migration, 'op', fake_op)
+    monkeypatch.setattr(
+        migration.sa,
+        'inspect',
+        lambda _conn: _FakeInspector(tables={'monitoring_traces', 'monitoring_spans'}),
+    )
+
+    migration.upgrade()
+
+    assert fake_op.created_tables == []
+    assert fake_op.created_indexes == []
+
+
+def test_downgrade_drops_spans_before_traces(monkeypatch):
+    migration = _migration_module()
+    fake_op = _FakeOp()
+
+    monkeypatch.setattr(migration, 'op', fake_op)
+    monkeypatch.setattr(
+        migration.sa,
+        'inspect',
+        lambda _conn: _FakeInspector(tables={'monitoring_traces', 'monitoring_spans'}),
+    )
+
+    migration.downgrade()
+
+    assert fake_op.dropped_tables == ['monitoring_spans', 'monitoring_traces']