33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
import importlib
|
|
import os
|
|
import unittest
|
|
|
|
|
|
def _set_required_env() -> None:
|
|
os.environ["SCW_ACCESS_KEY"] = "key0"
|
|
os.environ["SCW_SECRET_KEY"] = "key1"
|
|
os.environ["SCW_DEFAULT_ORGANIZATION_ID"] = "key2"
|
|
os.environ["SCW_DEFAULT_PROJECT_ID"] = "key3"
|
|
os.environ["SCW_SERVERLESS"] = "key4"
|
|
|
|
|
|
class ApiClientModelSelectionTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
_set_required_env()
|
|
import migration_pipeline.api_client as api_client
|
|
|
|
self.api_client = importlib.reload(api_client)
|
|
|
|
def test_default_model_family_is_glm(self) -> None:
|
|
self.assertEqual(self.api_client.get_selected_model_family(), "glm")
|
|
self.assertEqual(self.api_client.select_model(10), "glm-5.2")
|
|
|
|
def test_switch_to_qwen_model(self) -> None:
|
|
self.api_client.set_model_family("qwen")
|
|
self.assertEqual(self.api_client.select_model(10), "qwen3.5-397b-a17b")
|
|
|
|
def test_invalid_model_family_raises(self) -> None:
|
|
with self.assertRaises(ValueError):
|
|
self.api_client.set_model_family("invalid")
|
|
|