274 lines
9.3 KiB
Python
274 lines
9.3 KiB
Python
"""GSM 配置快照（进程内读路径）相关测试。
|
||
|
||
去 Ray 后，快照不再进 Ray Object Store：``_publish_snapshot`` 直接在进程内构建
|
||
并持有一个不可变 ``GSMSnapshot`` 对象，``current_config_ref`` 返回 ``(version, snapshot)``。
|
||
读端 ``fetch_snapshot`` 用版本号做进程内缓存失效。
|
||
|
||
主要验证：
|
||
|
||
- ``_build_snapshot`` 正确从 6 类内存状态打包配置
|
||
- ``_publish_snapshot`` 让 version 单调递增并刷新快照对象
|
||
- 写入路径（add_individual / add_provider_wrap / 等）会自动发布新快照
|
||
- ``fetch_snapshot`` 客户端：版本号一致时走本地缓存，不一致时重拉
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
import pytest
|
||
|
||
from kilostar.core.global_state_machine.gsm_snapshot import (
|
||
GSMSnapshot,
|
||
fetch_snapshot,
|
||
reset_local_cache,
|
||
)
|
||
|
||
|
||
# ─── GSM 集成（直接构造 plain 对象） ────────────────────────────────
|
||
|
||
|
||
@pytest.fixture
|
||
def gsm_instance():
|
||
from kilostar.core.global_state_machine.global_state_machine import (
|
||
GlobalStateMachine,
|
||
)
|
||
from kilostar.core.global_state_machine.individual_manager import (
|
||
GlobalIndividualManager,
|
||
)
|
||
from kilostar.core.global_state_machine.provider_manager import ProviderManager
|
||
from kilostar.core.global_state_machine.skill_manager import GlobalSkillManager
|
||
from kilostar.core.global_state_machine.tool_manager import GlobalToolManager
|
||
|
||
obj = GlobalStateMachine.__new__(GlobalStateMachine)
|
||
obj._global_provider_manager = ProviderManager(postgres=None)
|
||
obj._global_tool_manager = GlobalToolManager()
|
||
obj._global_skill_manager = GlobalSkillManager()
|
||
obj._global_individual_manager = GlobalIndividualManager()
|
||
obj._mcp_servers = {}
|
||
obj._tool_configs = {}
|
||
obj._custom_toolsets = {}
|
||
obj._config_version = 0
|
||
obj._current_ref = None
|
||
obj.postgres_database = MagicMock()
|
||
return obj
|
||
|
||
|
||
def test_build_snapshot_picks_up_all_six_categories(gsm_instance):
|
||
"""_build_snapshot 应正确从 GSM 内存的 6 类数据打包。"""
|
||
from kilostar.core.global_state_machine.model_provider.base_provider import (
|
||
Provider,
|
||
)
|
||
|
||
gsm_instance._global_provider_manager.provider_register["p1"] = Provider(
|
||
provider_title="p1",
|
||
provider_url="http://x",
|
||
provider_apikey="k",
|
||
provider_models=[],
|
||
provider_type="openai",
|
||
)
|
||
gsm_instance._global_individual_manager._individuals["a1"] = {"agent_id": "a1"}
|
||
gsm_instance._mcp_servers["s1"] = {"server_id": "s1"}
|
||
gsm_instance._tool_configs["t1"] = {"key": "v"}
|
||
gsm_instance._custom_toolsets["ts1"] = {"toolset_id": "ts1"}
|
||
|
||
snap = gsm_instance._build_snapshot()
|
||
|
||
assert "p1" in snap.providers
|
||
assert "a1" in snap.individuals
|
||
assert "s1" in snap.mcp_servers
|
||
assert "t1" in snap.tool_configs
|
||
assert "ts1" in snap.custom_toolsets
|
||
|
||
|
||
def test_build_snapshot_exposes_system_tools_by_scope(gsm_instance):
|
||
"""系统工具按 scope 分桶的工具名清单要随快照发布出去（客户端重建 toolset 用）。"""
|
||
tm = gsm_instance._global_tool_manager
|
||
|
||
def _f1():
|
||
return "f1"
|
||
|
||
def _f2():
|
||
return "f2"
|
||
|
||
tm._tool_funcs.clear()
|
||
tm._tool_funcs["default"]["file_reader"] = _f1
|
||
tm._tool_funcs["control_node"]["approval"] = _f2
|
||
|
||
snap = gsm_instance._build_snapshot()
|
||
assert snap.system_tools_by_scope.get("default") == ["file_reader"]
|
||
assert snap.system_tools_by_scope.get("control_node") == ["approval"]
|
||
assert set(snap.tool_funcs.keys()) == {"file_reader", "approval"}
|
||
|
||
|
||
def test_publish_snapshot_increments_version(gsm_instance):
|
||
assert gsm_instance._config_version == 0
|
||
assert gsm_instance._current_ref is None
|
||
|
||
gsm_instance._publish_snapshot()
|
||
v1 = gsm_instance._config_version
|
||
ref1 = gsm_instance._current_ref
|
||
assert v1 == 1
|
||
assert ref1 is not None
|
||
|
||
gsm_instance._publish_snapshot()
|
||
assert gsm_instance._config_version == 2
|
||
assert gsm_instance._current_ref is not ref1 # 重新构建应是新对象
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_current_config_ref_lazy_publishes_when_empty(gsm_instance):
|
||
"""从未发布过快照时，current_config_ref 应自动发布一次而不是返回 None。"""
|
||
version, snap = await gsm_instance.current_config_ref()
|
||
assert version == 1
|
||
assert snap is not None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_current_version_is_lightweight(gsm_instance):
|
||
gsm_instance._publish_snapshot()
|
||
gsm_instance._publish_snapshot()
|
||
assert await gsm_instance.current_version() == 2
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_add_individual_publishes_new_snapshot(gsm_instance):
|
||
"""写入路径 add_individual 应自动 +1 version。"""
|
||
before = gsm_instance._config_version
|
||
await gsm_instance.add_individual("agent-x", {"model_id": "gpt-4o"})
|
||
after = gsm_instance._config_version
|
||
assert after == before + 1
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_add_provider_wrap_publishes_new_snapshot(gsm_instance):
|
||
"""add_provider_wrap 即便走 mock 适配器也应该最终发布一次新快照。"""
|
||
from kilostar.core.global_state_machine.model_provider.base_provider import (
|
||
Provider,
|
||
)
|
||
|
||
fake_provider = Provider(
|
||
provider_title="my-openai",
|
||
provider_url="http://x",
|
||
provider_apikey="k",
|
||
provider_models=[],
|
||
provider_type="openai",
|
||
)
|
||
gsm_instance._global_provider_manager.provider_mapper["openai"] = MagicMock()
|
||
gsm_instance._global_provider_manager.provider_mapper[
|
||
"openai"
|
||
].create_provider = AsyncMock(return_value=fake_provider)
|
||
gsm_instance.postgres_database.add_provider_db = AsyncMock()
|
||
|
||
before = gsm_instance._config_version
|
||
await gsm_instance.add_provider_wrap(
|
||
provider_type="openai",
|
||
provider_title="my-openai",
|
||
provider_url="http://x",
|
||
provider_apikey="k",
|
||
provider_owner="alice",
|
||
)
|
||
after = gsm_instance._config_version
|
||
assert after == before + 1
|
||
|
||
|
||
# ─── fetch_snapshot 客户端缓存 ────────────────────────────────────
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_fetch_snapshot_uses_local_cache_when_version_matches():
|
||
"""版本号一致时不再调 current_config_ref，直接返回本地缓存快照。"""
|
||
reset_local_cache()
|
||
snap = GSMSnapshot(version=5, providers={"p": MagicMock()})
|
||
|
||
fake_gsm = MagicMock()
|
||
fake_gsm.current_version = AsyncMock(return_value=5)
|
||
fake_gsm.current_config_ref = AsyncMock(
|
||
side_effect=AssertionError("不应触发：缓存版本一致时不应调 current_config_ref")
|
||
)
|
||
|
||
from kilostar.core.global_state_machine import gsm_snapshot as snap_mod
|
||
|
||
snap_mod._local_cache["version"] = 5
|
||
snap_mod._local_cache["snapshot"] = snap
|
||
|
||
result = await fetch_snapshot(gsm_actor=fake_gsm)
|
||
assert result is snap
|
||
fake_gsm.current_version.assert_awaited_once()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_fetch_snapshot_refetches_when_version_changes():
|
||
"""版本号变了应重新拉取新 snapshot 并更新缓存。"""
|
||
reset_local_cache()
|
||
new_snap = GSMSnapshot(version=10)
|
||
|
||
fake_gsm = MagicMock()
|
||
fake_gsm.current_version = AsyncMock(return_value=10)
|
||
fake_gsm.current_config_ref = AsyncMock(return_value=(10, new_snap))
|
||
|
||
result = await fetch_snapshot(gsm_actor=fake_gsm)
|
||
assert result is new_snap
|
||
fake_gsm.current_config_ref.assert_awaited_once()
|
||
|
||
from kilostar.core.global_state_machine import gsm_snapshot as snap_mod
|
||
|
||
assert snap_mod._local_cache["version"] == 10
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_fetch_snapshot_use_cache_false_skips_cache():
|
||
"""``use_cache=False`` 直接走 current_config_ref，不读本地缓存。"""
|
||
reset_local_cache()
|
||
fresh = GSMSnapshot(version=1)
|
||
|
||
fake_gsm = MagicMock()
|
||
fake_gsm.current_config_ref = AsyncMock(return_value=(1, fresh))
|
||
|
||
result = await fetch_snapshot(gsm_actor=fake_gsm, use_cache=False)
|
||
assert result is fresh
|
||
|
||
|
||
# ─── build_tools_for_scope 客户端 helper ────────────────────────
|
||
|
||
|
||
def test_build_tools_for_scope_assembles_system_and_custom():
|
||
"""客户端按 snapshot 的 custom_toolsets + all_funcs 展开为扁平 callable 列表。"""
|
||
from kilostar.core.global_state_machine.gsm_snapshot import (
|
||
build_tools_for_scope,
|
||
)
|
||
|
||
def _sys_default():
|
||
return "d"
|
||
|
||
def _sys_scope():
|
||
return "s"
|
||
|
||
def _tp_a():
|
||
return "a"
|
||
|
||
snap = GSMSnapshot(
|
||
all_funcs={"sys_default": _sys_default, "sys_scope": _sys_scope, "tp_a": _tp_a},
|
||
custom_toolsets={
|
||
"system_basic": {
|
||
"toolset_id": "system_basic",
|
||
"tools": ["sys_default", "sys_scope"],
|
||
},
|
||
"grp": {"toolset_id": "grp", "tools": ["tp_a"]},
|
||
},
|
||
)
|
||
|
||
result = build_tools_for_scope(snap, "control_node")
|
||
assert len(result) == 3
|
||
assert result == [_sys_default, _sys_scope, _tp_a]
|
||
|
||
|
||
def test_build_tools_for_scope_skips_empty_buckets():
|
||
"""没有工具的 scope 返回空列表。"""
|
||
from kilostar.core.global_state_machine.gsm_snapshot import (
|
||
build_tools_for_scope,
|
||
)
|
||
|
||
snap = GSMSnapshot(all_funcs={}, custom_toolsets={})
|
||
assert build_tools_for_scope(snap, "control_node") == []
|