177 lines
6.9 KiB
Python
177 lines
6.9 KiB
Python
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
|
|
from services.account_service import AccountService
|
|
from services.model_service import ModelRoute, ModelUnavailableError
|
|
from services.protocol import (
|
|
anthropic_v1_messages,
|
|
conversation,
|
|
openai_v1_chat_complete,
|
|
openai_v1_response,
|
|
)
|
|
from services.storage.json_storage import JSONStorageBackend
|
|
|
|
|
|
class TextAccountRoutingTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp_dir.cleanup)
|
|
self.service = AccountService(
|
|
JSONStorageBackend(Path(self.temp_dir.name) / "accounts.json")
|
|
)
|
|
self.service.add_account_items(
|
|
[
|
|
{"access_token": "free", "type": "free", "status": "正常"},
|
|
{"access_token": "plus", "type": "Plus", "status": "正常"},
|
|
{"access_token": "pro", "type": "Pro", "status": "正常"},
|
|
{"access_token": "pro-disabled", "type": "Pro", "status": "禁用"},
|
|
]
|
|
)
|
|
self.service.refresh_access_token = lambda token, **_kwargs: token
|
|
|
|
def test_explicit_model_selects_only_advertising_account_type(self) -> None:
|
|
route = ModelRoute(account_types=frozenset({"Pro"}), allow_anonymous=False)
|
|
with mock.patch(
|
|
"services.model_service.model_catalog_service.route_for_model",
|
|
return_value=route,
|
|
):
|
|
token = self.service.get_text_access_token(model="pro-only")
|
|
|
|
self.assertEqual(token, "pro")
|
|
|
|
def test_auto_model_keeps_existing_unfiltered_rotation(self) -> None:
|
|
with mock.patch(
|
|
"services.model_service.model_catalog_service.route_for_model",
|
|
side_effect=AssertionError("auto must not load the model catalog"),
|
|
):
|
|
tokens = {
|
|
self.service.get_text_access_token(model="auto"),
|
|
self.service.get_text_access_token(model="auto"),
|
|
self.service.get_text_access_token(model="auto"),
|
|
}
|
|
|
|
self.assertEqual(tokens, {"free", "plus", "pro"})
|
|
|
|
def test_anonymous_model_uses_anonymous_backend(self) -> None:
|
|
route = ModelRoute(account_types=frozenset(), allow_anonymous=True)
|
|
with mock.patch(
|
|
"services.model_service.model_catalog_service.route_for_model",
|
|
return_value=route,
|
|
):
|
|
token = self.service.get_text_access_token(model="anon-only")
|
|
|
|
self.assertEqual(token, "")
|
|
|
|
def test_model_without_eligible_account_fails_closed(self) -> None:
|
|
route = ModelRoute(account_types=frozenset({"Team"}), allow_anonymous=False)
|
|
with mock.patch(
|
|
"services.model_service.model_catalog_service.route_for_model",
|
|
return_value=route,
|
|
), self.assertRaisesRegex(ModelUnavailableError, "team-only"):
|
|
self.service.get_text_access_token(model="team-only")
|
|
|
|
|
|
class TextProtocolRoutingTests(unittest.TestCase):
|
|
def test_text_backend_passes_requested_model_to_account_selector(self) -> None:
|
|
backend = mock.Mock()
|
|
with (
|
|
mock.patch.object(
|
|
conversation.account_service,
|
|
"get_text_access_token",
|
|
return_value="pro",
|
|
) as selector,
|
|
mock.patch.object(conversation, "OpenAIBackendAPI", return_value=backend),
|
|
):
|
|
result = conversation.text_backend("pro-only")
|
|
|
|
self.assertIs(result, backend)
|
|
selector.assert_called_once_with(model="pro-only")
|
|
|
|
def test_chat_completions_passes_requested_model_to_text_backend(self) -> None:
|
|
body = {
|
|
"model": "pro-chat",
|
|
"messages": [{"role": "user", "content": "route chat"}],
|
|
}
|
|
with (
|
|
mock.patch.object(openai_v1_chat_complete, "text_backend", return_value=object()) as backend,
|
|
mock.patch.object(openai_v1_chat_complete, "collect_text", return_value="ok"),
|
|
):
|
|
openai_v1_chat_complete.handle(body)
|
|
|
|
backend.assert_called_once_with("pro-chat")
|
|
|
|
def test_responses_passes_requested_model_to_text_backend(self) -> None:
|
|
body = {"model": "pro-response", "input": "route response"}
|
|
with (
|
|
mock.patch.object(openai_v1_response, "text_backend", return_value=object()) as backend,
|
|
mock.patch.object(openai_v1_response, "stream_text_deltas", return_value=iter(["ok"])),
|
|
):
|
|
openai_v1_response.handle(body)
|
|
|
|
backend.assert_called_once_with("pro-response")
|
|
|
|
def test_anthropic_messages_passes_requested_model_to_account_selector(self) -> None:
|
|
with (
|
|
mock.patch.object(
|
|
anthropic_v1_messages.account_service,
|
|
"get_text_access_token",
|
|
return_value="pro",
|
|
) as selector,
|
|
mock.patch.object(anthropic_v1_messages, "OpenAIBackendAPI"),
|
|
):
|
|
request = anthropic_v1_messages.message_request({
|
|
"model": "pro-anthropic",
|
|
"messages": [{"role": "user", "content": "route anthropic"}],
|
|
})
|
|
|
|
self.assertEqual(request.model, "pro-anthropic")
|
|
selector.assert_called_once_with(model="pro-anthropic")
|
|
|
|
def test_invalid_token_retry_keeps_requested_model_filter(self) -> None:
|
|
initial_backend = SimpleNamespace(access_token="bad")
|
|
request = conversation.ConversationRequest(
|
|
model="pro-only",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
)
|
|
|
|
def fake_events(backend, **_kwargs):
|
|
if backend.access_token == "bad":
|
|
raise RuntimeError("token_invalidated")
|
|
yield {"type": "conversation.delta", "delta": "ok"}
|
|
|
|
with (
|
|
mock.patch.object(conversation, "OpenAIBackendAPI", side_effect=lambda access_token: SimpleNamespace(
|
|
access_token=access_token,
|
|
close=lambda: None,
|
|
)),
|
|
mock.patch.object(conversation, "conversation_events", side_effect=fake_events),
|
|
mock.patch.object(
|
|
conversation.account_service,
|
|
"refresh_access_token",
|
|
return_value="bad",
|
|
),
|
|
mock.patch.object(conversation.account_service, "remove_invalid_token"),
|
|
mock.patch.object(
|
|
conversation.account_service,
|
|
"get_text_access_token",
|
|
return_value="pro",
|
|
) as selector,
|
|
mock.patch.object(conversation.account_service, "mark_text_used"),
|
|
):
|
|
result = list(conversation.stream_text_deltas(initial_backend, request))
|
|
|
|
self.assertEqual(result, ["ok"])
|
|
selector.assert_called_once_with(
|
|
excluded_tokens={"bad"},
|
|
model="pro-only",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|