65 lines
2.3 KiB
Python
65 lines
2.3 KiB
Python
"""Tests for the legacy model configuration migration."""
|
|
|
|
from agenteval.config import get_settings
|
|
from agenteval.storage.db import ModelConfigDB, ScenarioDB, ScenarioModelBindingDB
|
|
from cryptography.fernet import Fernet
|
|
from sqlmodel import select
|
|
|
|
from scripts.migrate_model_configs import migrate
|
|
|
|
|
|
def _legacy_scenario() -> ScenarioDB:
|
|
scenario = ScenarioDB(id="legacy-scenario", name="legacy")
|
|
scenario.set_llm_config(
|
|
{"api_url": "https://models.example.com/v1/chat/completions", "api_key": "secret", "model": "chat"}
|
|
)
|
|
scenario.set_cases(
|
|
[
|
|
{
|
|
"id": "case-1",
|
|
"messages": ["hello"],
|
|
"eval_rules": [
|
|
{
|
|
"type": "llm_score",
|
|
"params": {
|
|
"api_url": "https://models.example.com/v1/chat/completions",
|
|
"api_key": "secret",
|
|
"model": "chat",
|
|
"criteria": "准确",
|
|
},
|
|
}
|
|
],
|
|
}
|
|
]
|
|
)
|
|
return scenario
|
|
|
|
|
|
def test_migration_dry_run_and_apply_are_secret_safe(db_session, monkeypatch):
|
|
scenario = _legacy_scenario()
|
|
db_session.add(scenario)
|
|
db_session.commit()
|
|
monkeypatch.setattr(get_settings(), "secret_key", Fernet.generate_key().decode("ascii"))
|
|
|
|
preview = migrate(db_session, apply=False)
|
|
assert preview == {"scenarios": 1, "bindings": 2, "configs": 1}
|
|
assert db_session.exec(select(ModelConfigDB)).all() == []
|
|
|
|
result = migrate(db_session, apply=True)
|
|
assert result == preview
|
|
configs = db_session.exec(select(ModelConfigDB)).all()
|
|
bindings = db_session.exec(select(ScenarioModelBindingDB)).all()
|
|
assert len(configs) == 1
|
|
assert len(bindings) == 2
|
|
assert configs[0].api_key_encrypted != "secret"
|
|
|
|
refreshed = db_session.get(ScenarioDB, scenario.id)
|
|
assert refreshed.get_llm_config() is None
|
|
rule_params = refreshed.get_cases()[0]["eval_rules"][0]["params"]
|
|
assert rule_params == {"criteria": "准确"}
|
|
|
|
rerun = migrate(db_session, apply=True)
|
|
assert rerun["bindings"] == 0
|
|
assert len(db_session.exec(select(ModelConfigDB)).all()) == 1
|
|
assert len(db_session.exec(select(ScenarioModelBindingDB)).all()) == 2
|