521 lines
18 KiB
Python
521 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from yuxi.channel.plugins.protocol import (
|
|
BindingConversationRef,
|
|
BindingRoute,
|
|
InboundMessage,
|
|
SessionConversationRef,
|
|
)
|
|
from yuxi.channel.routing.router import BindingRouter, _BindingContext
|
|
|
|
|
|
@pytest.fixture
|
|
def plugin():
|
|
mock = MagicMock()
|
|
mock.compile_binding = MagicMock(return_value=None)
|
|
mock.match_binding = MagicMock(return_value=False)
|
|
return mock
|
|
|
|
|
|
@pytest.fixture
|
|
def route_cache():
|
|
return MagicMock()
|
|
|
|
|
|
@pytest.fixture
|
|
def router(route_cache):
|
|
return BindingRouter(route_cache=route_cache)
|
|
|
|
|
|
@pytest.fixture
|
|
def inbound():
|
|
return InboundMessage(
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
session_key="feishu:acc1:dm:user1",
|
|
sender_id="user1",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def session_ref():
|
|
return SessionConversationRef(
|
|
session_key="feishu:acc1:dm:user1",
|
|
chat_type="dm",
|
|
channel_sender_id="user1",
|
|
)
|
|
|
|
|
|
class TestBindingRouterHelpers:
|
|
def test_build_cache_key(self, router):
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
assert router._build_cache_key(ctx) == "feishu:acc1:feishu:acc1:dm:user1"
|
|
|
|
def test_compute_rule_hash_is_stable(self, router):
|
|
rule = {"agent_id": "a1", "channel_type": "feishu"}
|
|
payload = json.dumps(rule, sort_keys=True, ensure_ascii=False, default=str)
|
|
expected = hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
|
assert router._compute_rule_hash(rule) == expected
|
|
|
|
def test_parse_session_key_full(self, router):
|
|
assert router._parse_session_key("feishu:acc1:dm:user1") == ("feishu", "acc1")
|
|
|
|
def test_parse_session_key_empty(self, router):
|
|
assert router._parse_session_key("") == (None, None)
|
|
|
|
def test_parse_session_key_single_part(self, router):
|
|
assert router._parse_session_key("feishu") == ("feishu", None)
|
|
|
|
def test_parse_session_key_two_parts(self, router):
|
|
assert router._parse_session_key("feishu:acc1") == ("feishu", "acc1")
|
|
|
|
def test_parse_session_key_many_parts(self, router):
|
|
assert router._parse_session_key("feishu:acc1:dm:user1:thread") == (
|
|
"feishu",
|
|
"acc1",
|
|
)
|
|
|
|
def test_make_context_from_config(self, router, plugin, session_ref):
|
|
config = {"channel_type": "feishu", "account_id": "acc1"}
|
|
ctx = router._make_context(config, plugin, session_ref)
|
|
assert ctx.channel_type == "feishu"
|
|
assert ctx.account_id == "acc1"
|
|
assert ctx.session_key == "feishu:acc1:dm:user1"
|
|
|
|
def test_make_context_from_session_key(self, router, plugin, session_ref):
|
|
config = {}
|
|
ctx = router._make_context(config, plugin, session_ref)
|
|
assert ctx.channel_type == "feishu"
|
|
assert ctx.account_id == "acc1"
|
|
|
|
def test_make_context_raises_when_unresolvable(self, router, plugin):
|
|
ref = SessionConversationRef(session_key="single")
|
|
with pytest.raises(ValueError, match="Unable to resolve"):
|
|
router._make_context({}, plugin, ref)
|
|
|
|
|
|
class TestResolveConfiguredBinding:
|
|
async def test_plugin_binding_takes_precedence(self, router, plugin, inbound):
|
|
plugin.compile_binding.return_value = BindingConversationRef(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
)
|
|
plugin.match_binding.return_value = True
|
|
config = {
|
|
"bindings": [
|
|
{"agent_id": "plugin_agent", "custom": True},
|
|
]
|
|
}
|
|
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
route = await router._resolve_configured_binding(ctx, config, plugin, inbound)
|
|
|
|
assert route.agent_id == "plugin_agent"
|
|
assert route.matched_by == "binding.plugin"
|
|
plugin.compile_binding.assert_called_once_with(config, inbound)
|
|
plugin.match_binding.assert_called_once_with(config, config["bindings"][0], inbound)
|
|
|
|
async def test_plugin_binding_not_used_when_compile_returns_none(self, router, plugin, inbound):
|
|
plugin.compile_binding.return_value = None
|
|
config = {
|
|
"bindings": [
|
|
{
|
|
"agent_id": "tier_agent",
|
|
"channel_type": "feishu",
|
|
"account_id": "acc1",
|
|
}
|
|
]
|
|
}
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
route = await router._resolve_configured_binding(ctx, config, plugin, inbound)
|
|
|
|
assert route.agent_id == "tier_agent"
|
|
assert route.matched_by == "binding.channel_account"
|
|
plugin.match_binding.assert_not_called()
|
|
|
|
async def test_no_bindings_returns_none(self, router, plugin, inbound):
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
route = await router._resolve_configured_binding(ctx, {}, plugin, inbound)
|
|
assert route is None
|
|
|
|
async def test_session_key_tier_matches(self, router, plugin, inbound):
|
|
config = {
|
|
"bindings": [
|
|
{
|
|
"agent_id": "agent1",
|
|
"session_key": "feishu:acc1:dm:user1",
|
|
}
|
|
]
|
|
}
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
route = await router._resolve_configured_binding(ctx, config, plugin, inbound)
|
|
assert route.agent_id == "agent1"
|
|
assert route.matched_by == "binding.session_key"
|
|
|
|
async def test_channel_account_chat_peer_tier_matches(self, router, plugin, inbound):
|
|
config = {
|
|
"bindings": [
|
|
{
|
|
"agent_id": "agent1",
|
|
"channel_type": "feishu",
|
|
"account_id": "acc1",
|
|
"chat_type": "dm",
|
|
"peer_id": "user1",
|
|
}
|
|
]
|
|
}
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
route = await router._resolve_configured_binding(ctx, config, plugin, inbound)
|
|
assert route.matched_by == "binding.channel_account_chat_peer"
|
|
|
|
async def test_channel_account_peer_tier_matches(self, router, plugin, inbound):
|
|
config = {
|
|
"bindings": [
|
|
{
|
|
"agent_id": "agent1",
|
|
"channel_type": "feishu",
|
|
"account_id": "acc1",
|
|
"peer_id": "user1",
|
|
}
|
|
]
|
|
}
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
route = await router._resolve_configured_binding(ctx, config, plugin, inbound)
|
|
assert route.matched_by == "binding.channel_account_peer"
|
|
|
|
async def test_channel_account_chat_tier_matches(self, router, plugin, inbound):
|
|
config = {
|
|
"bindings": [
|
|
{
|
|
"agent_id": "agent1",
|
|
"channel_type": "feishu",
|
|
"account_id": "acc1",
|
|
"chat_type": "dm",
|
|
}
|
|
]
|
|
}
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
route = await router._resolve_configured_binding(ctx, config, plugin, inbound)
|
|
assert route.matched_by == "binding.channel_account_chat"
|
|
|
|
async def test_channel_chat_peer_tier_matches(self, router, plugin, inbound):
|
|
config = {
|
|
"bindings": [
|
|
{
|
|
"agent_id": "agent1",
|
|
"channel_type": "feishu",
|
|
"chat_type": "dm",
|
|
"peer_id": "user1",
|
|
}
|
|
]
|
|
}
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
route = await router._resolve_configured_binding(ctx, config, plugin, inbound)
|
|
assert route.matched_by == "binding.channel_chat_peer"
|
|
|
|
async def test_session_key_regex_tier_matches(self, router, plugin, inbound):
|
|
config = {
|
|
"bindings": [
|
|
{
|
|
"agent_id": "agent1",
|
|
"session_key_regex": "feishu:acc1:.*",
|
|
}
|
|
]
|
|
}
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
route = await router._resolve_configured_binding(ctx, config, plugin, inbound)
|
|
assert route.matched_by == "binding.session_key_regex"
|
|
|
|
async def test_channel_account_tier_matches(self, router, plugin, inbound):
|
|
config = {
|
|
"bindings": [
|
|
{
|
|
"agent_id": "agent1",
|
|
"channel_type": "feishu",
|
|
"account_id": "acc1",
|
|
}
|
|
]
|
|
}
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
route = await router._resolve_configured_binding(ctx, config, plugin, inbound)
|
|
assert route.matched_by == "binding.channel_account"
|
|
|
|
async def test_channel_type_tier_matches(self, router, plugin, inbound):
|
|
config = {
|
|
"bindings": [
|
|
{
|
|
"agent_id": "agent1",
|
|
"channel_type": "feishu",
|
|
}
|
|
]
|
|
}
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
route = await router._resolve_configured_binding(ctx, config, plugin, inbound)
|
|
assert route.matched_by == "binding.channel_type"
|
|
|
|
async def test_peer_id_falls_back_to_context(self, router, plugin):
|
|
inbound_no_peer = InboundMessage(
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id=None,
|
|
session_key="feishu:acc1:dm:user1",
|
|
)
|
|
config = {
|
|
"bindings": [
|
|
{
|
|
"agent_id": "agent1",
|
|
"channel_type": "feishu",
|
|
"account_id": "acc1",
|
|
"peer_id": "user1",
|
|
}
|
|
]
|
|
}
|
|
ctx = _BindingContext(
|
|
session_key="feishu:acc1:dm:user1",
|
|
channel_type="feishu",
|
|
account_id="acc1",
|
|
chat_type="dm",
|
|
peer_id="user1",
|
|
)
|
|
route = await router._resolve_configured_binding(ctx, config, plugin, inbound_no_peer)
|
|
assert route.matched_by == "binding.channel_account_peer"
|
|
|
|
|
|
class TestResolveRuntime:
|
|
async def test_cache_hit_returns_cached_route(self, router, route_cache, plugin, session_ref):
|
|
cached = BindingRoute(agent_id="cached", session_key="k", matched_by="cache")
|
|
route_cache.get = AsyncMock(return_value=cached)
|
|
route_cache.set = AsyncMock()
|
|
|
|
route = await router.resolve_runtime(
|
|
{"channel_type": "feishu", "account_id": "acc1"},
|
|
plugin,
|
|
session_ref,
|
|
MagicMock(),
|
|
)
|
|
|
|
assert route is cached
|
|
route_cache.get.assert_awaited_once()
|
|
route_cache.set.assert_not_called()
|
|
|
|
async def test_configured_binding_caches_and_returns(self, router, route_cache, plugin, session_ref):
|
|
route_cache.get = AsyncMock(return_value=None)
|
|
route_cache.set = AsyncMock()
|
|
plugin.compile_binding.return_value = None
|
|
config = {
|
|
"channel_type": "feishu",
|
|
"account_id": "acc1",
|
|
"bindings": [{"agent_id": "bound", "channel_type": "feishu", "account_id": "acc1"}],
|
|
}
|
|
|
|
route = await router.resolve_runtime(config, plugin, session_ref, MagicMock())
|
|
|
|
assert route.agent_id == "bound"
|
|
assert route.matched_by == "binding.channel_account"
|
|
route_cache.set.assert_awaited_once()
|
|
|
|
async def test_runtime_binding_caches_and_returns(self, router, route_cache, plugin, session_ref):
|
|
route_cache.get = AsyncMock(return_value=None)
|
|
route_cache.set = AsyncMock()
|
|
plugin.compile_binding.return_value = None
|
|
|
|
binding = MagicMock()
|
|
binding.agent_id = "runtime_agent"
|
|
binding.binding_rule_hash = "hash1"
|
|
|
|
with patch("yuxi.channel.routing.router.ChannelBindingRepository") as mock_repo_cls:
|
|
mock_repo = MagicMock()
|
|
mock_repo.find_runtime_bindings = AsyncMock(return_value=[binding])
|
|
mock_repo_cls.return_value = mock_repo
|
|
|
|
route = await router.resolve_runtime(
|
|
{"channel_type": "feishu", "account_id": "acc1"},
|
|
plugin,
|
|
session_ref,
|
|
MagicMock(),
|
|
)
|
|
|
|
assert route.agent_id == "runtime_agent"
|
|
assert route.matched_by == "runtime.binding"
|
|
route_cache.set.assert_awaited_once()
|
|
|
|
async def test_default_route_uses_config_default_agent(self, router, route_cache, plugin, session_ref):
|
|
route_cache.get = AsyncMock(return_value=None)
|
|
route_cache.set = AsyncMock()
|
|
plugin.compile_binding.return_value = None
|
|
config = {
|
|
"channel_type": "feishu",
|
|
"account_id": "acc1",
|
|
"default_agent_id": "default_agent",
|
|
}
|
|
|
|
with patch("yuxi.channel.routing.router.ChannelBindingRepository") as mock_repo_cls:
|
|
mock_repo = MagicMock()
|
|
mock_repo.find_runtime_bindings = AsyncMock(return_value=[])
|
|
mock_repo_cls.return_value = mock_repo
|
|
|
|
route = await router.resolve_runtime(config, plugin, session_ref, MagicMock())
|
|
|
|
assert route.agent_id == "default_agent"
|
|
assert route.matched_by == "default"
|
|
|
|
async def test_default_route_without_agent_id_not_cached(self, router, route_cache, plugin, session_ref):
|
|
route_cache.get = AsyncMock(return_value=None)
|
|
route_cache.set = AsyncMock()
|
|
plugin.compile_binding.return_value = None
|
|
config = {"channel_type": "feishu", "account_id": "acc1"}
|
|
|
|
with patch("yuxi.channel.routing.router.ChannelBindingRepository") as mock_repo_cls:
|
|
mock_repo = MagicMock()
|
|
mock_repo.find_runtime_bindings = AsyncMock(return_value=[])
|
|
mock_repo_cls.return_value = mock_repo
|
|
|
|
route = await router.resolve_runtime(config, plugin, session_ref, MagicMock())
|
|
|
|
assert route.agent_id is None
|
|
route_cache.set.assert_not_called()
|
|
|
|
|
|
class TestResolve:
|
|
async def test_cache_hit_returns_cached_route(self, router, route_cache, plugin, inbound):
|
|
cached = BindingRoute(agent_id="cached", session_key="k", matched_by="cache")
|
|
route_cache.get = AsyncMock(return_value=cached)
|
|
route_cache.set = AsyncMock()
|
|
|
|
session = MagicMock()
|
|
session.session_key = "feishu:acc1:dm:user1"
|
|
session.channel_type = "feishu"
|
|
session.account_id = "acc1"
|
|
session.chat_type = "dm"
|
|
session.channel_sender_id = "user1"
|
|
session.conversation_id = 1
|
|
|
|
route = await router.resolve(session, {"default_agent_id": "x"}, plugin, inbound, MagicMock())
|
|
|
|
assert route is cached
|
|
route_cache.set.assert_not_called()
|
|
|
|
async def test_falls_back_to_conversation_agent_when_no_default(self, router, route_cache, plugin, inbound):
|
|
route_cache.get = AsyncMock(return_value=None)
|
|
route_cache.set = AsyncMock()
|
|
plugin.compile_binding.return_value = None
|
|
|
|
conversation = MagicMock()
|
|
conversation.agent_id = "conv_agent"
|
|
|
|
with (
|
|
patch("yuxi.channel.routing.router.ConversationRepository") as mock_conv_repo_cls,
|
|
patch("yuxi.channel.routing.router.ChannelBindingRepository") as mock_binding_repo_cls,
|
|
):
|
|
mock_conv_repo = MagicMock()
|
|
mock_conv_repo.get_conversation_by_id = AsyncMock(return_value=conversation)
|
|
mock_conv_repo_cls.return_value = mock_conv_repo
|
|
|
|
mock_binding_repo = MagicMock()
|
|
mock_binding_repo.find_runtime_bindings = AsyncMock(return_value=[])
|
|
mock_binding_repo_cls.return_value = mock_binding_repo
|
|
|
|
session = MagicMock()
|
|
session.session_key = "feishu:acc1:dm:user1"
|
|
session.channel_type = "feishu"
|
|
session.account_id = "acc1"
|
|
session.chat_type = "dm"
|
|
session.channel_sender_id = "user1"
|
|
session.conversation_id = 1
|
|
|
|
route = await router.resolve(session, {}, plugin, inbound, MagicMock())
|
|
|
|
assert route.agent_id == "conv_agent"
|
|
assert route.matched_by == "default"
|
|
|
|
|
|
class TestInvalidate:
|
|
async def test_invalidate_channel_cache_delegates_to_cache(self, router, route_cache):
|
|
route_cache.invalidate_by_channel = AsyncMock(return_value=3)
|
|
assert await router.invalidate_channel_cache("feishu") == 3
|
|
|
|
async def test_invalidate_account_cache_delegates_to_cache(self, router, route_cache):
|
|
route_cache.invalidate_by_account = AsyncMock(return_value=2)
|
|
assert await router.invalidate_account_cache("feishu", "acc1") == 2
|