Route open models through ZenMux with OpenRouter fallback

2026-10-11 18:40:45 -07:00
parent 9ebc213d71
commit cd50972093
4 changed files with 93 additions and 5 deletions
@@ -54,6 +54,35 @@ class ConfigTests(unittest.TestCase):
        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']}])
+        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)
+

 if __name__ == '__main__':
    unittest.main()