126 lines
7.2 KiB
Python
126 lines
7.2 KiB
Python
import unittest
|
|
import yaml
|
|
from pathlib import Path
|
|
from tempfile import TemporaryDirectory
|
|
from render_config import build_config, oauth_ready
|
|
|
|
|
|
class ConfigTests(unittest.TestCase):
|
|
catalog = {
|
|
'zenmux': {'claude-sonnet-latest': 'anthropic/claude-sonnet-5.5'},
|
|
'openrouter': {'glm-latest': 'z-ai/glm-5.3-prime'},
|
|
'google_ai_studio': {'gemini-latest': 'gemini-test'},
|
|
'google_vertex': {'gemini-vertex': 'gemini-vertex-test'},
|
|
'chatgpt': {'codex-latest': 'gpt-test'},
|
|
'xai': {'grok-latest': 'grok-test'},
|
|
}
|
|
|
|
def test_missing_credentials_exclude_routes(self):
|
|
self.assertEqual(build_config(self.catalog, {}, lambda _: False)['model_list'], [])
|
|
|
|
def test_namespaces_preserve_every_legacy_route(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'
|
|
old = build_config(catalog, env, lambda _: True, namespaced=False)
|
|
new = build_config(catalog, env, lambda _: True)
|
|
routes = {r['model_name']: r for r in new['model_list']}
|
|
self.assertEqual(len(routes), len(new['model_list']))
|
|
for route in old['model_list']:
|
|
self.assertEqual(route, routes[route['model_name']])
|
|
for old_name, new_name in {
|
|
'claude-sonnet-latest': 'zenmux/claude-sonnet-latest',
|
|
'claude-opus-latest-payg': 'zenmux-payg/claude-opus-latest',
|
|
'grok-latest': 'xai/grok-latest', 'gpt-sol-latest': 'codex/gpt-sol-latest',
|
|
'gpt-work-sol': 'codex/gpt-6-sol', 'gpt-review-5.5': 'codex/gpt-5.5',
|
|
'gemini-transcribe-personal': 'aistudio/gemini-transcribe-personal',
|
|
'gemini-embedding-2-studio': 'aistudio/gemini-embedding-2',
|
|
'google-rerank-fast': 'vertex/rerank-fast-latest',
|
|
}.items():
|
|
self.assertEqual(routes[old_name]['litellm_params'], routes[new_name]['litellm_params'])
|
|
|
|
def test_env_ui_login_disabled(self):
|
|
settings = build_config(self.catalog, {}, lambda _: False)['general_settings']
|
|
self.assertTrue(settings['disable_env_credential_login'])
|
|
|
|
def test_subscription_preserves_vendor_prefix(self):
|
|
routes = build_config(self.catalog, {'ZENMUX_SUBSCRIPTION_API_KEY': 'secret'}, lambda _: False)['model_list']
|
|
self.assertEqual(routes[0]['litellm_params']['model'], 'anthropic/anthropic/claude-sonnet-5.5')
|
|
self.assertNotIn('secret', str(routes))
|
|
|
|
def test_payg_separate_alias(self):
|
|
routes = build_config(self.catalog, {'ZENMUX_PAYG_API_KEY': 'secret'}, lambda _: False)['model_list']
|
|
self.assertEqual(routes[0]['model_name'], 'claude-sonnet-latest-payg')
|
|
|
|
def test_oauth_requires_persistent_file(self):
|
|
routes = build_config(self.catalog, {}, lambda path: str(path) == '/tokens/xai/auth.json')['model_list']
|
|
self.assertEqual(len(routes), 2)
|
|
self.assertTrue(routes[0]['litellm_params']['use_xai_oauth'])
|
|
self.assertNotIn('api_key', routes[0]['litellm_params'])
|
|
|
|
def test_pending_oauth_is_not_ready(self):
|
|
with TemporaryDirectory() as folder:
|
|
path = Path(folder) / 'auth.json'
|
|
self.assertFalse(oauth_ready(path))
|
|
path.write_text('{"device_code_requested_at": 1}')
|
|
self.assertFalse(oauth_ready(path))
|
|
path.write_text('{"access_token": "fixture-only"}')
|
|
self.assertTrue(oauth_ready(path))
|
|
|
|
def test_vertex_requires_explicit_enable_and_project(self):
|
|
for env in ({'VERTEXAI_PROJECT': 'test-project'}, {'VERTEXAI_ENABLED': 'true'}):
|
|
self.assertEqual(build_config(self.catalog, env, lambda _: False)['model_list'], [])
|
|
routes = build_config(self.catalog, {'VERTEXAI_ENABLED': 'true',
|
|
'VERTEXAI_PROJECT': 'test-project'}, lambda _: False)['model_list']
|
|
self.assertEqual(routes[0]['litellm_params']['model'], 'vertex_ai/gemini-vertex-test')
|
|
self.assertEqual(routes[0]['litellm_params']['vertex_project'], 'os.environ/VERTEXAI_PROJECT')
|
|
self.assertNotIn('api_key', routes[0]['litellm_params'])
|
|
|
|
def test_open_models_subscription_first(self):
|
|
catalog = {**self.catalog, 'zenmux_open_models': self.catalog['openrouter']}
|
|
config = build_config(catalog, {'ZENMUX_SUBSCRIPTION_API_KEY': 'a',
|
|
'OPENROUTER_API_KEY': 'b', 'ZENMUX_PAYG_API_KEY': 'c'}, lambda _: False)
|
|
routes = {r['model_name']: r for r in config['model_list']}
|
|
self.assertEqual(routes['glm-latest']['litellm_params']['model'], 'openai/z-ai/glm-5.3-prime')
|
|
self.assertEqual(routes['glm-latest-openrouter']['litellm_params']['model'], 'openrouter/z-ai/glm-5.3-prime')
|
|
self.assertEqual(config['router_settings']['fallbacks'], [
|
|
{'glm-latest': ['glm-latest-openrouter']},
|
|
{'auto/glm-latest': ['openrouter/glm-latest']}])
|
|
self.assertEqual(routes['auto/glm-latest']['litellm_params'], routes['glm-latest']['litellm_params'])
|
|
self.assertEqual(routes['zenmux/glm-latest']['litellm_params'], routes['glm-latest']['litellm_params'])
|
|
self.assertNotIn('zenmux/glm-latest', str(config['router_settings']['fallbacks']))
|
|
self.assertEqual(routes['glm-latest']['model_info']['cooldown_time'], 60)
|
|
self.assertNotIn('glm-latest-payg', routes)
|
|
|
|
def test_openrouter_only_remains_usable(self):
|
|
catalog = {**self.catalog, 'zenmux_open_models': self.catalog['openrouter']}
|
|
config = build_config(catalog, {'OPENROUTER_API_KEY': 'b'}, lambda _: False)
|
|
self.assertEqual(config['model_list'][0]['litellm_params']['model'], 'openrouter/z-ai/glm-5.3-prime')
|
|
self.assertEqual(config['router_settings']['fallbacks'], [])
|
|
|
|
def test_zenmux_only_no_broken_fallback(self):
|
|
catalog = {**self.catalog, 'zenmux_open_models': self.catalog['openrouter']}
|
|
config = build_config(catalog, {'ZENMUX_SUBSCRIPTION_API_KEY': 'a'}, lambda _: False)
|
|
self.assertIn('glm-latest', [r['model_name'] for r in config['model_list']])
|
|
self.assertEqual(config['router_settings']['fallbacks'], [])
|
|
|
|
def test_mismatched_model_rejected(self):
|
|
catalog = {**self.catalog, 'zenmux_open_models': {'glm-latest': 'different-model'}}
|
|
with self.assertRaises(ValueError):
|
|
build_config(catalog, {'ZENMUX_SUBSCRIPTION_API_KEY': 'a',
|
|
'OPENROUTER_API_KEY': 'b'}, lambda _: False)
|
|
|
|
def test_zenmux_chat_keeps_existing_data_destination(self):
|
|
catalog = {**self.catalog, 'zenmux_chat': {'gemini-flash-zenmux': 'google/gemini-3.8-flash'}}
|
|
config = build_config(catalog, {'ZENMUX_SUBSCRIPTION_API_KEY': 'a'}, lambda _: False)
|
|
route = next(r for r in config['model_list'] if r['model_name'] == 'gemini-flash-zenmux')
|
|
self.assertEqual(route['litellm_params']['model'], 'openai/google/gemini-3.8-flash')
|
|
self.assertEqual(route['litellm_params']['api_base'], 'https://zenmux.ai/api/v1')
|
|
self.assertNotIn('gemini-flash-zenmux-payg', [r['model_name'] for r in config['model_list']])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|