mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-17 19:24:19 +00:00
huanghuoguoguo
88 lines
2.5 KiB
Python
88 lines
2.5 KiB
Python
"""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']
|