224 lines
8.4 KiB
Python
224 lines
8.4 KiB
Python
"""U15 — migrate_api_keys_to_secrets 顶层函数单元测试。
|
||
|
||
测试 ``src/agentkit/llm/migration.py`` 中的 ``migrate_api_keys_to_secrets(config_path)``：
|
||
读取 ``agentkit.yaml``，把每个 LLM provider 的 plaintext ``api_key`` 迁移到
|
||
``SecretsStore`` 加密存储，并把更新后的 providers 段写回 YAML。
|
||
|
||
覆盖场景：
|
||
1. Happy path — plaintext → secrets_store，YAML 写回，report 标记 migrated
|
||
2. Idempotent — 已迁移配置再次调用返回 skipped，不重复加密
|
||
3. Empty config — 无 providers 时返回 {}，YAML 仍可读
|
||
4. Source missing — config_path 不存在时抛 FileNotFoundError（函数未捕获）
|
||
5. Secrets store write failure — set_secret 抛异常时 status=failed，plaintext 保留（事务性）
|
||
6. Dual source re-migrate — api_key_source=dual 时重新迁移覆盖旧 encrypted 值
|
||
|
||
隔离策略：
|
||
- 用 ``tmp_path`` 隔离 YAML 文件 I/O
|
||
- patch ``agentkit.channels.secrets.SecretsStore`` 注入可控实例（避免 env 依赖）
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
||
import pytest
|
||
import yaml
|
||
|
||
from agentkit.channels.secrets import SecretsStore
|
||
from agentkit.llm.migration import migrate_api_keys_to_secrets
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# 测试辅助
|
||
# ----------------------------------------------------------------------
|
||
|
||
|
||
def _write_config(path: Path, providers: dict[str, dict[str, object]]) -> None:
|
||
"""写一个最小的 agentkit.yaml（仅 llm.providers 段）。"""
|
||
data: dict[str, dict[str, object]] = {"llm": {"providers": providers}}
|
||
path.write_text(
|
||
yaml.safe_dump(data, allow_unicode=True, sort_keys=False),
|
||
encoding="utf-8",
|
||
)
|
||
|
||
|
||
def _make_real_store() -> SecretsStore:
|
||
"""构造一个真实可用的 SecretsStore（注入固定 master_key，不依赖 env）。"""
|
||
return SecretsStore(master_key=b"x" * 32)
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# 1. Happy path
|
||
# ----------------------------------------------------------------------
|
||
|
||
|
||
def test_happy_path_migrates_plaintext_to_secrets_store(tmp_path: Path) -> None:
|
||
"""plaintext api_key → secrets_store；YAML 写回，report 标记 migrated。"""
|
||
config_path = tmp_path / "agentkit.yaml"
|
||
_write_config(
|
||
config_path,
|
||
{
|
||
"openai": {
|
||
"api_key": "sk-secret",
|
||
"base_url": "https://api.openai.com/v1",
|
||
"type": "openai",
|
||
}
|
||
},
|
||
)
|
||
|
||
with patch("agentkit.channels.secrets.SecretsStore", return_value=_make_real_store()):
|
||
report = migrate_api_keys_to_secrets(config_path)
|
||
|
||
assert report == {"openai": {"status": "migrated", "source": "secrets_store"}}
|
||
|
||
# YAML 应已更新：plaintext 清空，encrypted 写入，source 标记
|
||
raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
||
pconf = raw["llm"]["providers"]["openai"]
|
||
assert pconf["api_key"] == ""
|
||
assert pconf["api_key_encrypted"] is not None
|
||
assert pconf["api_key_source"] == "secrets_store"
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# 2. Idempotent
|
||
# ----------------------------------------------------------------------
|
||
|
||
|
||
def test_idempotent_second_call_returns_skipped(tmp_path: Path) -> None:
|
||
"""已迁移的配置再次调用 → skipped，不重复加密。"""
|
||
config_path = tmp_path / "agentkit.yaml"
|
||
_write_config(
|
||
config_path,
|
||
{
|
||
"openai": {
|
||
"api_key": "sk-secret",
|
||
"base_url": "https://api.openai.com/v1",
|
||
"type": "openai",
|
||
}
|
||
},
|
||
)
|
||
|
||
with patch("agentkit.channels.secrets.SecretsStore", return_value=_make_real_store()):
|
||
first_report = migrate_api_keys_to_secrets(config_path)
|
||
second_report = migrate_api_keys_to_secrets(config_path)
|
||
|
||
assert first_report == {"openai": {"status": "migrated", "source": "secrets_store"}}
|
||
assert second_report == {"openai": {"status": "skipped", "source": "secrets_store"}}
|
||
|
||
# 第二次调用后 YAML 仍一致（plaintext 仍为空）
|
||
raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
||
pconf = raw["llm"]["providers"]["openai"]
|
||
assert pconf["api_key"] == ""
|
||
assert pconf["api_key_source"] == "secrets_store"
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# 3. Empty config
|
||
# ----------------------------------------------------------------------
|
||
|
||
|
||
def test_empty_config_returns_empty_report(tmp_path: Path) -> None:
|
||
"""无 providers 时返回 {}，YAML 仍可读且 providers 段为空。"""
|
||
config_path = tmp_path / "agentkit.yaml"
|
||
config_path.write_text("llm:\n providers: {}\n", encoding="utf-8")
|
||
|
||
with patch("agentkit.channels.secrets.SecretsStore", return_value=_make_real_store()):
|
||
report = migrate_api_keys_to_secrets(config_path)
|
||
|
||
assert report == {}
|
||
|
||
raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
||
assert raw["llm"]["providers"] == {}
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# 4. Source missing
|
||
# ----------------------------------------------------------------------
|
||
|
||
|
||
def test_source_missing_raises_file_not_found(tmp_path: Path) -> None:
|
||
"""config_path 不存在时抛 FileNotFoundError（函数未捕获，向上传播）。"""
|
||
missing = tmp_path / "nonexistent.yaml"
|
||
with pytest.raises(FileNotFoundError):
|
||
migrate_api_keys_to_secrets(missing)
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# 5. Secrets store write failure
|
||
# ----------------------------------------------------------------------
|
||
|
||
|
||
def test_secrets_store_write_failure_marks_failed_and_retains_plaintext(
|
||
tmp_path: Path,
|
||
) -> None:
|
||
"""set_secret 抛异常 → status=failed，plaintext 保留（事务性，不部分迁移）。"""
|
||
config_path = tmp_path / "agentkit.yaml"
|
||
_write_config(
|
||
config_path,
|
||
{
|
||
"openai": {
|
||
"api_key": "sk-secret",
|
||
"base_url": "",
|
||
"type": "openai",
|
||
}
|
||
},
|
||
)
|
||
|
||
# mock store：set_secret 抛异常（get_secret 不会被调用到）
|
||
fake_store = MagicMock(spec=SecretsStore)
|
||
fake_store.set_secret = AsyncMock(side_effect=RuntimeError("redis down"))
|
||
|
||
with patch("agentkit.channels.secrets.SecretsStore", return_value=fake_store):
|
||
report = migrate_api_keys_to_secrets(config_path)
|
||
|
||
assert report == {"openai": {"status": "failed", "source": "plaintext", "error": "redis down"}}
|
||
|
||
# plaintext 应保留（未清空），source 仍为 plaintext
|
||
raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
||
pconf = raw["llm"]["providers"]["openai"]
|
||
assert pconf["api_key"] == "sk-secret"
|
||
assert pconf["api_key_source"] == "plaintext"
|
||
|
||
# set_secret 被调用一次（验证读 get_secret 未执行）
|
||
fake_store.set_secret.assert_awaited_once()
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# 6. Dual source re-migrate (secret overwrite path)
|
||
# ----------------------------------------------------------------------
|
||
|
||
|
||
def test_dual_source_re_migrates_and_overwrites_encrypted(tmp_path: Path) -> None:
|
||
"""api_key_source=dual（已有 encrypted 但 plaintext 保留）→ 重新迁移覆盖旧值。"""
|
||
stale_encrypted = (
|
||
'{"key":"llm:provider:openai:api_key","value":"old","nonce":"n","salt":"s",'
|
||
'"key_id":"default","created_at":"","updated_at":""}'
|
||
)
|
||
config_path = tmp_path / "agentkit.yaml"
|
||
_write_config(
|
||
config_path,
|
||
{
|
||
"openai": {
|
||
"api_key": "sk-secret",
|
||
"base_url": "",
|
||
"type": "openai",
|
||
"api_key_encrypted": stale_encrypted,
|
||
"api_key_source": "dual",
|
||
}
|
||
},
|
||
)
|
||
|
||
with patch("agentkit.channels.secrets.SecretsStore", return_value=_make_real_store()):
|
||
report = migrate_api_keys_to_secrets(config_path)
|
||
|
||
assert report == {"openai": {"status": "migrated", "source": "secrets_store"}}
|
||
|
||
raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
||
pconf = raw["llm"]["providers"]["openai"]
|
||
assert pconf["api_key"] == ""
|
||
assert pconf["api_key_source"] == "secrets_store"
|
||
# encrypted 应被新值覆盖（不等于原来的 stale 占位）
|
||
assert pconf["api_key_encrypted"] != stale_encrypted
|
||
assert pconf["api_key_encrypted"] is not None
|