shaguabijia-app-server/tests/test_deepseek_price_migration.py

from __future__ import annotations

import importlib.util
from datetime import datetime
from pathlib import Path

import sqlalchemy as sa


def _load_migration():
    path = (
        Path(__file__).parents[1]
        / "alembic"
        / "versions"
        / "deepseek_v4_flash_price.py"
    )
    spec = importlib.util.spec_from_file_location("deepseek_v4_flash_price", path)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def test_migration_follows_comparison_status_normalization():
    migration = _load_migration()

    assert migration.down_revision == "comparison_below_min_success"


def test_migration_adds_price_and_corrects_only_mispriced_snapshot(monkeypatch):
    migration = _load_migration()
    engine = sa.create_engine("sqlite://")
    metadata = sa.MetaData()
    app_config = sa.Table(
        "app_config",
        metadata,
        sa.Column("key", sa.String(64), primary_key=True),
        sa.Column("value", sa.JSON, nullable=False),
        sa.Column("updated_at", sa.DateTime),
    )
    comparison = sa.Table(
        "comparison_record",
        metadata,
        sa.Column("id", sa.Integer, primary_key=True),
        sa.Column("llm_calls", sa.JSON),
        sa.Column("llm_cost_yuan", sa.Float),
        sa.Column("llm_price_snapshot", sa.JSON),
    )
    metadata.create_all(engine)

    calls = [
        {
            "model": "deepseek-v4-flash",
            "error": None,
            "usage": {"prompt_tokens": 12031, "completion_tokens": 125},
        },
        {
            "model": "qwen3.5-flash",
            "error": None,
            "usage": {"prompt_tokens": 2416, "completion_tokens": 119},
        },
    ]
    snapshot = {
        "mode": "per_model",
        "prices": {
            "deepseek-v4-flash": {
                "input_per_1m": 3.0,
                "output_per_1m": 15.0,
                "_source": "default",
            },
            "qwen3.5-flash": {
                "input_per_1m": 0.2,
                "output_per_1m": 2.0,
                "_source": "per_model",
            },
        },
    }

    with engine.begin() as conn:
        original_updated_at = datetime(2026, 7, 13, 18, 20, 10)
        conn.execute(
            app_config.insert().values(
                key="llm_token_price",
                value={
                    "per_model": {
                        "qwen3.5-flash": {
                            "input_per_1m": 0.2,
                            "output_per_1m": 2.0,
                        }
                    },
                    "default": {"input_per_1m": 3.0, "output_per_1m": 15.0},
                },
                updated_at=original_updated_at,
            )
        )
        conn.execute(
            comparison.insert().values(
                id=1,
                llm_calls=calls,
                llm_cost_yuan=0.038689,
                llm_price_snapshot=snapshot,
            )
        )
        monkeypatch.setattr(migration.op, "get_bind", lambda: conn)

        migration.upgrade()

        config = conn.execute(
            sa.select(app_config.c.value).where(
                app_config.c.key == "llm_token_price"
            )
        ).scalar_one()
        assert config["per_model"]["deepseek-v4-flash"] == {
            "input_per_1m": 1.0,
            "output_per_1m": 2.0,
        }
        assert conn.execute(
            sa.select(app_config.c.updated_at).where(
                app_config.c.key == "llm_token_price"
            )
        ).scalar_one() == original_updated_at
        corrected = conn.execute(sa.select(comparison)).mappings().one()
        assert corrected["llm_cost_yuan"] == 0.013002
        assert corrected["llm_price_snapshot"]["prices"]["deepseek-v4-flash"] == {
            "input_per_1m": 1.0,
            "output_per_1m": 2.0,
            "_source": "per_model",
        }
        assert corrected["llm_price_snapshot"]["pricing_correction"] == (
            "deepseek_v4_flash_price"
        )

        migration.downgrade()

        reverted_config = conn.execute(
            sa.select(app_config.c.value).where(
                app_config.c.key == "llm_token_price"
            )
        ).scalar_one()
        assert "deepseek-v4-flash" not in reverted_config["per_model"]
        reverted = conn.execute(sa.select(comparison)).mappings().one()
        assert reverted["llm_cost_yuan"] == 0.038689
        assert "pricing_correction" not in reverted["llm_price_snapshot"]


def test_downgrade_preserves_price_that_existed_before_upgrade(monkeypatch):
    migration = _load_migration()
    engine = sa.create_engine("sqlite://")
    metadata = sa.MetaData()
    app_config = sa.Table(
        "app_config",
        metadata,
        sa.Column("key", sa.String(64), primary_key=True),
        sa.Column("value", sa.JSON, nullable=False),
        sa.Column("updated_at", sa.DateTime),
    )
    sa.Table(
        "comparison_record",
        metadata,
        sa.Column("id", sa.Integer, primary_key=True),
        sa.Column("llm_calls", sa.JSON),
        sa.Column("llm_cost_yuan", sa.Float),
        sa.Column("llm_price_snapshot", sa.JSON),
    )
    metadata.create_all(engine)

    explicit_price = {"input_per_1m": 1.0, "output_per_1m": 2.0}
    with engine.begin() as conn:
        conn.execute(
            app_config.insert().values(
                key="llm_token_price",
                value={
                    "per_model": {"deepseek-v4-flash": explicit_price},
                    "default": {"input_per_1m": 3.0, "output_per_1m": 15.0},
                },
            )
        )
        monkeypatch.setattr(migration.op, "get_bind", lambda: conn)

        migration.upgrade()
        migration.downgrade()

        config = conn.execute(sa.select(app_config.c.value)).scalar_one()
        assert config["per_model"]["deepseek-v4-flash"] == explicit_price