50 lines
2.8 KiB
Python
50 lines
2.8 KiB
Python
import unittest
|
|
from pathlib import Path
|
|
import yaml
|
|
from render_config import build_config
|
|
|
|
|
|
class PolicyTests(unittest.TestCase):
|
|
def build(self, **env):
|
|
catalog = yaml.safe_load(Path(__file__).with_name('models.yaml').read_text())
|
|
return build_config(catalog, env, lambda _: True)
|
|
|
|
def test_subscription_is_ignored_even_if_old_env_still_exists(self):
|
|
c = self.build(ZENMUX_SUBSCRIPTION_API_KEY='retired-test')
|
|
self.assertNotIn('ZENMUX_SUBSCRIPTION_API_KEY', yaml.safe_dump(c))
|
|
self.assertFalse(any('zenmux.ai' in r['litellm_params'].get('api_base', '') for r in c['model_list']))
|
|
|
|
def test_openrouter_primary_and_byok_last_for_all_supported_aliases(self):
|
|
c = self.build(ZENMUX_SUBSCRIPTION_API_KEY='retired', ZENMUX_PAYG_API_KEY='test', OPENROUTER_API_KEY='test')
|
|
routes = {r['model_name']: r['litellm_params'] for r in c['model_list']}
|
|
groups = {k:v for row in c['router_settings']['fallbacks'] for k,v in row.items()}
|
|
for alias in ['claude-sonnet-latest', 'claude-opus-latest', 'claude-haiku-latest',
|
|
'deepseek-flash-latest', 'glm-latest', 'kimi-latest', 'mimo-flash-latest']:
|
|
for name in [alias, 'auto/'+alias, 'zenmux/'+alias, 'openrouter/'+alias]:
|
|
self.assertEqual(routes[name]['api_key'], 'os.environ/OPENROUTER_API_KEY')
|
|
self.assertEqual(groups[name][-1], 'zenmux-payg/'+alias)
|
|
self.assertEqual(routes['zenmux-payg/'+alias]['api_key'], 'os.environ/ZENMUX_PAYG_API_KEY')
|
|
for source, targets in groups.items():
|
|
self.assertTrue(all(target in routes for target in targets))
|
|
self.assertNotEqual(routes[source].get('api_key'), 'os.environ/ZENMUX_PAYG_API_KEY')
|
|
self.assertEqual(len(routes), len(c['model_list']))
|
|
|
|
def test_no_byok_key_no_dangling_fallback_and_no_subscription(self):
|
|
c = self.build(OPENROUTER_API_KEY='test', ZENMUX_SUBSCRIPTION_API_KEY='retired')
|
|
self.assertEqual(c['router_settings']['fallbacks'], [])
|
|
self.assertNotIn('ZENMUX_SUBSCRIPTION_API_KEY', yaml.safe_dump(c))
|
|
|
|
def test_oauth_official_fallback_stays_unchanged(self):
|
|
c = self.build(OPENAI_API_KEY='test', OPENROUTER_API_KEY='test', ZENMUX_PAYG_API_KEY='test')
|
|
groups = {k:v for row in c['router_settings']['fallbacks'] for k,v in row.items()}
|
|
self.assertEqual(groups['codex/gpt-sol-latest'], ['openai/gpt-sol-latest'])
|
|
|
|
def test_google_primary_before_gemini_byok(self):
|
|
c = self.build(GOOGLE_AI_STUDIO_API_KEY='test', ZENMUX_PAYG_API_KEY='test')
|
|
routes = {r['model_name']: r['litellm_params'] for r in c['model_list']}
|
|
self.assertEqual(routes['zenmux/gemini-flash-latest']['api_key'], 'os.environ/GOOGLE_AI_STUDIO_API_KEY')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|