52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
"""Verify model metadata migration preserves existing configurations."""
|
|
|
|
import importlib
|
|
|
|
import sqlalchemy as sa
|
|
from alembic.migration import MigrationContext
|
|
from alembic.operations import Operations
|
|
|
|
|
|
def test_model_metadata_migration_adds_safe_defaults(tmp_path, monkeypatch):
|
|
engine = sa.create_engine(f"sqlite:///{tmp_path / 'model_metadata.db'}")
|
|
metadata = sa.MetaData()
|
|
legacy_table = sa.Table(
|
|
"model_configs",
|
|
metadata,
|
|
sa.Column("id", sa.String(), primary_key=True),
|
|
sa.Column("name", sa.String(), nullable=False),
|
|
)
|
|
metadata.create_all(engine)
|
|
|
|
with engine.begin() as connection:
|
|
connection.execute(legacy_table.insert().values(id="existing", name="现有模型"))
|
|
operations = Operations(MigrationContext.configure(connection))
|
|
migration = importlib.import_module("migrations.versions.a64b2f8c9d10_add_model_metadata")
|
|
monkeypatch.setattr(migration, "op", operations)
|
|
migration.upgrade()
|
|
|
|
columns = {column["name"] for column in sa.inspect(connection).get_columns("model_configs")}
|
|
assert {
|
|
"vendor_name",
|
|
"input_modalities",
|
|
"output_modalities",
|
|
"context_window",
|
|
"supports_tool_calling",
|
|
"documentation_url",
|
|
} <= columns
|
|
|
|
row = connection.execute(
|
|
sa.text(
|
|
"SELECT name, vendor_name, input_modalities, output_modalities, supports_streaming "
|
|
"FROM model_configs WHERE id = :id"
|
|
),
|
|
{"id": "existing"},
|
|
).mappings().one()
|
|
assert row == {
|
|
"name": "现有模型",
|
|
"vendor_name": "",
|
|
"input_modalities": '["text"]',
|
|
"output_modalities": '["text"]',
|
|
"supports_streaming": 0,
|
|
}
|