12a4d9584e
Stabilize cross-protocol ask-mode/streaming behavior and reduce session-reuse branch collisions, then add focused docs/tests for multimodal normalization and pool/stats/config paths. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
194 lines
7.4 KiB
Python
194 lines
7.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
# app.lingma_pool imports auto_login; tests here don't execute Playwright paths.
|
|
# Stub module import so test environments without playwright can import pool code.
|
|
_playwright = types.ModuleType("playwright")
|
|
_playwright_async = types.ModuleType("playwright.async_api")
|
|
|
|
|
|
class _StubPlaywrightTimeoutError(Exception):
|
|
pass
|
|
|
|
|
|
async def _stub_async_playwright():
|
|
raise RuntimeError("playwright is stubbed in unit tests")
|
|
|
|
|
|
_playwright_async.TimeoutError = _StubPlaywrightTimeoutError
|
|
_playwright_async.async_playwright = _stub_async_playwright
|
|
sys.modules.setdefault("playwright", _playwright)
|
|
sys.modules.setdefault("playwright.async_api", _playwright_async)
|
|
|
|
from app.config import _parse_accounts, load_settings
|
|
from app.lingma_pool import LingmaPool
|
|
from app.stats import StatsCollector, estimate_tokens
|
|
|
|
|
|
def _affinity_key_for_bucket(pool_size: int, bucket_index: int) -> str:
|
|
for i in range(20000):
|
|
key = f"k-{i}"
|
|
if abs(hash(key)) % pool_size == bucket_index:
|
|
return key
|
|
raise RuntimeError("failed to find affinity key")
|
|
|
|
|
|
class _FakeInstance:
|
|
def __init__(self, idx: int, *, healthy: bool, in_flight: int):
|
|
self.name = f"inst-{idx}"
|
|
self.cfg = SimpleNamespace(index=idx)
|
|
self._healthy = healthy
|
|
self.in_flight = in_flight
|
|
|
|
@property
|
|
def healthy(self) -> bool:
|
|
return self._healthy
|
|
|
|
|
|
class LingmaPoolRoutingTests(unittest.TestCase):
|
|
def test_pool_pick_prefers_healthy_affinity_bucket(self) -> None:
|
|
inst0 = _FakeInstance(0, healthy=True, in_flight=0)
|
|
inst1 = _FakeInstance(1, healthy=True, in_flight=9)
|
|
pool = LingmaPool([inst0, inst1])
|
|
|
|
key = _affinity_key_for_bucket(2, 1)
|
|
picked = pool.pick(affinity_key=key)
|
|
|
|
self.assertIs(picked, inst1)
|
|
|
|
def test_pool_pick_falls_back_to_least_in_flight_when_affinity_unhealthy(self) -> None:
|
|
inst0 = _FakeInstance(0, healthy=True, in_flight=1)
|
|
inst1 = _FakeInstance(1, healthy=False, in_flight=0)
|
|
inst2 = _FakeInstance(2, healthy=True, in_flight=1)
|
|
pool = LingmaPool([inst0, inst1, inst2])
|
|
|
|
key = _affinity_key_for_bucket(3, 1)
|
|
picked = pool.pick(affinity_key=key)
|
|
|
|
self.assertIs(picked, inst0)
|
|
|
|
def test_pool_pick_round_robin_when_all_unhealthy(self) -> None:
|
|
inst0 = _FakeInstance(0, healthy=False, in_flight=0)
|
|
inst1 = _FakeInstance(1, healthy=False, in_flight=0)
|
|
inst2 = _FakeInstance(2, healthy=False, in_flight=0)
|
|
pool = LingmaPool([inst0, inst1, inst2])
|
|
|
|
self.assertIs(pool.pick(), inst0)
|
|
self.assertIs(pool.pick(), inst1)
|
|
self.assertIs(pool.pick(), inst2)
|
|
self.assertIs(pool.pick(), inst0)
|
|
|
|
def test_pool_prometheus_lines_include_required_metrics(self) -> None:
|
|
inst0 = _FakeInstance(0, healthy=True, in_flight=2)
|
|
inst1 = _FakeInstance(1, healthy=False, in_flight=5)
|
|
pool = LingmaPool([inst0, inst1])
|
|
|
|
text = "\n".join(pool.prometheus_lines())
|
|
|
|
self.assertIn("# TYPE gateway_pool_instance_in_flight gauge", text)
|
|
self.assertIn("# TYPE gateway_pool_instance_ready gauge", text)
|
|
self.assertIn('gateway_pool_instance_in_flight{name="inst-0",idx="0"} 2', text)
|
|
self.assertIn('gateway_pool_instance_ready{name="inst-0",idx="0"} 1', text)
|
|
self.assertIn('gateway_pool_instance_ready{name="inst-1",idx="1"} 0', text)
|
|
|
|
|
|
class StatsCollectorTests(unittest.IsolatedAsyncioTestCase):
|
|
def test_estimate_tokens_empty_short_utf8(self) -> None:
|
|
self.assertEqual(estimate_tokens(""), 0)
|
|
self.assertGreaterEqual(estimate_tokens("a"), 1)
|
|
self.assertEqual(estimate_tokens("你好世界"), 3)
|
|
|
|
async def test_record_chat_updates_counters_and_clamps_negative_tokens(self) -> None:
|
|
s = StatsCollector()
|
|
|
|
await s.record_chat(stream=True, success=True, prompt_tokens=-3, completion_tokens=5)
|
|
await s.record_chat(stream=False, success=False, prompt_tokens=2, completion_tokens=-7)
|
|
snap = await s.snapshot()
|
|
|
|
self.assertEqual(snap["chat_requests_total"], 2)
|
|
self.assertEqual(snap["chat_requests_success"], 1)
|
|
self.assertEqual(snap["chat_requests_error"], 1)
|
|
self.assertEqual(snap["chat_stream_requests"], 1)
|
|
self.assertEqual(snap["chat_non_stream_requests"], 1)
|
|
self.assertEqual(snap["prompt_tokens_estimated_total"], 2)
|
|
self.assertEqual(snap["completion_tokens_estimated_total"], 5)
|
|
|
|
async def test_snapshot_and_prometheus_text_consistency(self) -> None:
|
|
s = StatsCollector()
|
|
|
|
await s.record_chat(stream=True, success=True, prompt_tokens=3, completion_tokens=4)
|
|
snap = await s.snapshot()
|
|
text = await s.prometheus_text()
|
|
|
|
self.assertEqual(snap["total_tokens_estimated"], 7)
|
|
self.assertIn("gateway_total_tokens_estimated 7", text)
|
|
self.assertIn("gateway_chat_requests_total 1", text)
|
|
self.assertTrue(text.endswith("\n"))
|
|
|
|
|
|
class ConfigParsingTests(unittest.TestCase):
|
|
def test_parse_accounts_accepts_json_csv_newline_formats(self) -> None:
|
|
raw_json = json.dumps([
|
|
{"username": "u1", "password": "p1"},
|
|
{"username": "u2", "password": "p2"},
|
|
])
|
|
parsed_json = _parse_accounts(raw_json)
|
|
self.assertEqual([a.username for a in parsed_json], ["u1", "u2"])
|
|
|
|
parsed_csv = _parse_accounts("u3:p3,u4:p4")
|
|
self.assertEqual([a.username for a in parsed_csv], ["u3", "u4"])
|
|
|
|
parsed_nl = _parse_accounts("u5:p5\nu6:p6")
|
|
self.assertEqual([a.username for a in parsed_nl], ["u5", "u6"])
|
|
|
|
def test_parse_accounts_allows_bundle_only_in_json(self) -> None:
|
|
raw = json.dumps([{"session_bundle": "abc"}])
|
|
parsed = _parse_accounts(raw)
|
|
|
|
self.assertEqual(len(parsed), 1)
|
|
self.assertEqual(parsed[0].username, "")
|
|
self.assertEqual(parsed[0].password, "")
|
|
self.assertEqual(parsed[0].session_bundle_b64, "abc")
|
|
|
|
def test_parse_accounts_csv_splits_only_first_colon(self) -> None:
|
|
parsed = _parse_accounts("u:p:with:colon")
|
|
|
|
self.assertEqual(len(parsed), 1)
|
|
self.assertEqual(parsed[0].username, "u")
|
|
self.assertEqual(parsed[0].password, "p:with:colon")
|
|
|
|
def test_load_settings_creates_bundle_only_account_without_credentials(self) -> None:
|
|
with patch.dict(os.environ, {"LINGMA_SESSION_BUNDLE": "abc"}, clear=True):
|
|
settings = load_settings()
|
|
|
|
self.assertEqual(len(settings.accounts), 1)
|
|
self.assertEqual(settings.accounts[0].username, "")
|
|
self.assertEqual(settings.accounts[0].password, "")
|
|
self.assertEqual(settings.accounts[0].session_bundle_b64, "abc")
|
|
|
|
def test_load_settings_invalid_instance_count_fallback(self) -> None:
|
|
with patch.dict(
|
|
os.environ,
|
|
{"LINGMA_ACCOUNTS": "u1:p1,u2:p2", "LINGMA_INSTANCE_COUNT": "not-a-number"},
|
|
clear=True,
|
|
):
|
|
settings_with_accounts = load_settings()
|
|
|
|
self.assertEqual(settings_with_accounts.instance_count, 2)
|
|
|
|
with patch.dict(os.environ, {"LINGMA_INSTANCE_COUNT": "not-a-number"}, clear=True):
|
|
settings_without_accounts = load_settings()
|
|
|
|
self.assertEqual(settings_without_accounts.instance_count, 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|