68 lines
3.9 KiB
Python
68 lines
3.9 KiB
Python
import unittest
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
|
|
from render_config import build_hk_config
|
|
|
|
|
|
class RoutingTests(unittest.TestCase):
|
|
def test_official_api_names_do_not_replace_hk_routes(self):
|
|
catalog = yaml.safe_load(Path(__file__).with_name('models.yaml').read_text())
|
|
env = {'ZENMUX_SUBSCRIPTION_API_KEY': 'synthetic', 'OPENAI_API_KEY': 'synthetic'}
|
|
config = build_hk_config(catalog, env)
|
|
routes = {r['model_name']: r['litellm_params'] for r in config['model_list']}
|
|
self.assertEqual(routes['openai/gpt-transcribe']['api_base'], 'https://api.openai.com/v1')
|
|
self.assertEqual(routes['zenmux/gpt-transcribe']['api_base'], 'https://zenmux.ai/api/v1')
|
|
self.assertEqual(routes['openai/gpt-sol-latest']['model'], 'openai/gpt-6.1-sol')
|
|
self.assertNotIn('codex/gpt-sol-latest', routes)
|
|
self.assertNotIn('synthetic', yaml.safe_dump(config))
|
|
|
|
def test_bailian_personal_uses_hk_private_vocabulary(self):
|
|
catalog = yaml.safe_load(Path(__file__).with_name('models.yaml').read_text())
|
|
routes = build_hk_config(catalog, {'BAILIAN_CN_API_KEY':'synthetic',
|
|
'BAILIAN_CN_API_BASE':'https://ws-test.cn-beijing.maas.aliyuncs.com/api/v1'})['model_list']
|
|
self.assertEqual(len(routes),2)
|
|
self.assertEqual(routes[1]['litellm_params']['bailian_vocabulary_path'],
|
|
'/credentials/personal-vocabulary.json')
|
|
|
|
def test_all_providers(self):
|
|
catalog = yaml.safe_load(Path(__file__).with_name('models.yaml').read_text())
|
|
env = {key: 'synthetic' for key in ['ZENMUX_SUBSCRIPTION_API_KEY',
|
|
'ZENMUX_PAYG_API_KEY', 'OPENROUTER_API_KEY', 'GOOGLE_AI_STUDIO_API_KEY',
|
|
'VERTEXAI_PROJECT']}
|
|
env['VERTEXAI_ENABLED'] = 'true'
|
|
result = build_hk_config(catalog, env)
|
|
routes = {r['model_name']: r['litellm_params'] for r in result['model_list']}
|
|
self.assertEqual(len(routes), len(result['model_list']))
|
|
self.assertEqual(routes['gemini-transcribe-personal']['model'], 'vertex_ai/gemini-3.5-transcribe')
|
|
self.assertIs(routes['gemini-transcribe-personal']['personal_dictation'], True)
|
|
self.assertEqual(routes['gemini-transcribe-personal']['vertex_location'], 'global')
|
|
self.assertNotIn('grok-latest', routes)
|
|
self.assertNotIn('gpt-sol-latest', routes)
|
|
self.assertEqual(routes['asr-gpt']['model'], 'openai/openai/gpt-transcribe')
|
|
self.assertEqual(routes['asr-qwen']['model'], 'openai/qwen/qwen-audio-3.1-asr-flash')
|
|
self.assertEqual(routes['asr-doubao']['model'], 'openai/bytedance/doubao-seed-asr-2.0')
|
|
self.assertEqual(routes['asr-mimo']['model'], 'openai/xiaomi/mimo-v2.5-asr')
|
|
for alias in ['asr-gpt', 'asr-qwen', 'asr-doubao', 'asr-mimo']:
|
|
self.assertNotIn('personal_dictation', routes[alias])
|
|
self.assertEqual(routes[alias]['api_key'], 'os.environ/ZENMUX_SUBSCRIPTION_API_KEY')
|
|
self.assertEqual(routes['gemini-transcribe']['model'], 'gemini/gemini-3.5-transcribe')
|
|
self.assertEqual(routes['gemini-3.6-flash']['model'], 'vertex_ai/gemini-3.8-flash')
|
|
self.assertEqual(len(result['router_settings']['fallbacks']), 12)
|
|
self.assertEqual(routes['vertex/gemini-transcribe-personal'], routes['gemini-transcribe-personal'])
|
|
self.assertNotIn('aistudio/gemini-transcribe-personal', routes)
|
|
self.assertEqual(routes['zenmux/gpt-transcribe'], routes['asr-gpt'])
|
|
self.assertEqual(result['general_settings'], {'master_key': 'os.environ/LITELLM_MASTER_KEY'})
|
|
self.assertNotIn('synthetic', yaml.safe_dump(result))
|
|
|
|
def test_missing_credentials(self):
|
|
catalog = yaml.safe_load(Path(__file__).with_name('models.yaml').read_text())
|
|
result = build_hk_config(catalog, {})
|
|
self.assertEqual(result['model_list'], [])
|
|
self.assertEqual(result['router_settings']['fallbacks'], [])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|