124 lines
5.8 KiB
Python
124 lines
5.8 KiB
Python
"""Regression coverage for bounded admission and inexpensive status reads."""
|
|
import json
|
|
import os
|
|
import sys
|
|
import threading
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'router'))
|
|
os.environ.setdefault('ROUTER_PROFILES_FILE', '')
|
|
import ai_profile_router as router
|
|
|
|
|
|
class CoordinationTests(unittest.TestCase):
|
|
def test_prompt_enhancer_accepts_plain_and_fenced_json(self):
|
|
plain = router._parse_prompt_enhancer_result(
|
|
'{"rewritten_prompt":"new scene","wh_ratio":"3:2"}')
|
|
fenced = router._parse_prompt_enhancer_result(
|
|
'```json\n{"rewritten_prompt":"portrait","wh_ratio":"3:4"}\n```')
|
|
self.assertEqual(plain['rewritten_prompt'], 'new scene')
|
|
self.assertEqual(fenced['wh_ratio'], '3:4')
|
|
|
|
def test_prompt_enhancer_rejects_missing_rewrite(self):
|
|
with self.assertRaisesRegex(RuntimeError, 'rewritten_prompt'):
|
|
router._parse_prompt_enhancer_result('{"wh_ratio":"1:1"}')
|
|
|
|
def test_image_data_url_resolves_volume_relative_reference(self):
|
|
with tempfile.TemporaryDirectory() as directory, \
|
|
patch.object(router, 'IMAGE_DIR', directory):
|
|
path = Path(directory) / '.edit-reference.ref'
|
|
path.write_bytes(b'\x89PNG\r\n\x1a\ncontent')
|
|
result = router._image_data_url(path.name)
|
|
self.assertTrue(result.startswith('data:image/png;base64,'))
|
|
|
|
def test_mode_uses_one_consistent_controller_snapshot(self):
|
|
with patch.object(router, 'PROFILE_CONTROL_URL', 'http://controller'), \
|
|
patch.object(router, '_profile_controller_request', return_value={
|
|
'music_worker': 'running', 'music_health': 'healthy'}) as request, \
|
|
patch.object(router.RUNTIME, 'load', return_value={}):
|
|
result = router.Handler._mode_payload()
|
|
self.assertEqual(result['music_worker'], 'running')
|
|
self.assertEqual(result['music_health'], 'healthy')
|
|
request.assert_called_once_with('GET', '/status', timeout=3)
|
|
|
|
def test_active_profile_does_not_request_optional_workers(self):
|
|
with patch.object(router, 'PROFILE_CONTROL_URL', 'http://controller'), \
|
|
patch.object(router, '_profile_controller_request',
|
|
return_value={'active_profile': 'medium'}) as request:
|
|
self.assertEqual(router.current_profile(), 'medium')
|
|
request.assert_called_once_with('GET', '/profiles/status', timeout=3)
|
|
|
|
def test_waiting_chat_times_out_without_upstream_work(self):
|
|
handler = object.__new__(router.Handler)
|
|
errors = []
|
|
handler._send_error = lambda *args: errors.append(args)
|
|
router.STATE.lock.acquire()
|
|
thread = None
|
|
try:
|
|
with patch.object(router, 'CHAT_WAIT_TIMEOUT', 0.02), \
|
|
patch.object(router, 'upstream_status') as upstream:
|
|
thread = threading.Thread(target=handler._chat_proxy,
|
|
args=(None, {'messages': []}, None))
|
|
thread.start()
|
|
thread.join(0.5)
|
|
self.assertFalse(thread.is_alive(), 'lock wait ignored timeout')
|
|
upstream.assert_not_called()
|
|
self.assertEqual(errors[0][0], 503)
|
|
self.assertEqual(errors[0][3], 'model_wait_timeout')
|
|
finally:
|
|
router.STATE.lock.release()
|
|
if thread:
|
|
thread.join(1)
|
|
|
|
def test_ultra_catalog_is_text_only(self):
|
|
with patch.object(router, 'current_profile', return_value='medium'):
|
|
catalog = router.Handler._llamacpp_models_payload()['data']
|
|
ultra = next(x for x in catalog if x['id'] == 'qwen-ultra')
|
|
self.assertEqual(ultra['architecture']['input_modalities'], ['text'])
|
|
|
|
def test_ultra_image_rejected_before_profile_switch(self):
|
|
handler = object.__new__(router.Handler)
|
|
errors = []
|
|
handler._send_error = lambda *args: errors.append(args)
|
|
data = {'messages': [{'role': 'user', 'content': [{'type': 'image_url',
|
|
'image_url': {'url': 'data:image/png;base64,iVBORw0KGgo='}}]}]}
|
|
with patch.object(router, 'switch_profile') as switch:
|
|
handler._chat_proxy(None, data, 'ultra')
|
|
switch.assert_not_called()
|
|
self.assertEqual(errors[0][0], 400)
|
|
|
|
def test_ready_chat_reuses_profile_readiness_result_and_releases_lease(self):
|
|
handler = object.__new__(router.Handler)
|
|
before = router.STATE.active_chats
|
|
available = router.STATE.qwen_unavailable
|
|
router.STATE.qwen_unavailable = False
|
|
handler._proxy = lambda body: self.assertEqual(
|
|
json.loads(body)['model'], 'qwen-medium')
|
|
try:
|
|
with patch.object(router, 'switch_profile', return_value={
|
|
'reachable': True, 'model': 'qwen-medium', 'ctx': 160000}), \
|
|
patch.object(router, 'upstream_status') as upstream:
|
|
handler._chat_proxy(None, {'messages': []}, 'medium')
|
|
upstream.assert_not_called()
|
|
self.assertEqual(router.STATE.active_chats, before)
|
|
finally:
|
|
router.STATE.qwen_unavailable = available
|
|
|
|
def test_startup_accepts_tested_mtp_context_overhead_without_restart(self):
|
|
with patch.object(router.RUNTIME, 'load', return_value={}), \
|
|
patch.object(router.RUNTIME, 'save'), \
|
|
patch.object(router, 'enforce_artifact_retention', return_value=[]), \
|
|
patch.object(router, 'current_profile', return_value='medium'), \
|
|
patch.object(router, 'upstream_status', return_value={
|
|
'reachable': True, 'model': 'qwen-medium', 'ctx': 160128}), \
|
|
patch.object(router, '_restore_qwen') as restore:
|
|
router._startup_reconcile()
|
|
restore.assert_not_called()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|