136 lines
4.1 KiB
Python
136 lines
4.1 KiB
Python
"""Tests for models module."""
|
|
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
# Add lib to path
|
|
sys.path.insert(0, str(Path(__file__).parent.parent / "scripts"))
|
|
|
|
from lib import models
|
|
|
|
|
|
class TestParseVersion(unittest.TestCase):
|
|
def test_simple_version(self):
|
|
result = models.parse_version("gpt-5")
|
|
self.assertEqual(result, (5,))
|
|
|
|
def test_minor_version(self):
|
|
result = models.parse_version("gpt-5.2")
|
|
self.assertEqual(result, (5, 2))
|
|
|
|
def test_patch_version(self):
|
|
result = models.parse_version("gpt-5.2.1")
|
|
self.assertEqual(result, (5, 2, 1))
|
|
|
|
def test_no_version(self):
|
|
result = models.parse_version("custom-model")
|
|
self.assertIsNone(result)
|
|
|
|
|
|
class TestIsMainlineOpenAIModel(unittest.TestCase):
|
|
def test_gpt5_is_mainline(self):
|
|
self.assertTrue(models.is_mainline_openai_model("gpt-5"))
|
|
|
|
def test_gpt52_is_mainline(self):
|
|
self.assertTrue(models.is_mainline_openai_model("gpt-5.2"))
|
|
|
|
def test_gpt5_mini_is_not_mainline(self):
|
|
self.assertFalse(models.is_mainline_openai_model("gpt-5-mini"))
|
|
|
|
def test_gpt4_is_not_mainline(self):
|
|
self.assertFalse(models.is_mainline_openai_model("gpt-4"))
|
|
|
|
|
|
class TestSelectOpenAIModel(unittest.TestCase):
|
|
def test_pinned_policy(self):
|
|
result = models.select_openai_model(
|
|
"fake-key",
|
|
policy="pinned",
|
|
pin="gpt-5.1"
|
|
)
|
|
self.assertEqual(result, "gpt-5.1")
|
|
|
|
def test_auto_with_mock_models(self):
|
|
mock_models = [
|
|
{"id": "gpt-5.2", "created": 1704067200},
|
|
{"id": "gpt-5.1", "created": 1701388800},
|
|
{"id": "gpt-5", "created": 1698710400},
|
|
]
|
|
result = models.select_openai_model(
|
|
"fake-key",
|
|
policy="auto",
|
|
mock_models=mock_models
|
|
)
|
|
self.assertEqual(result, "gpt-5.2")
|
|
|
|
def test_auto_filters_variants(self):
|
|
mock_models = [
|
|
{"id": "gpt-5.2", "created": 1704067200},
|
|
{"id": "gpt-5-mini", "created": 1704067200},
|
|
{"id": "gpt-5.1", "created": 1701388800},
|
|
]
|
|
result = models.select_openai_model(
|
|
"fake-key",
|
|
policy="auto",
|
|
mock_models=mock_models
|
|
)
|
|
self.assertEqual(result, "gpt-5.2")
|
|
|
|
|
|
class TestSelectXAIModel(unittest.TestCase):
|
|
def test_latest_policy(self):
|
|
result = models.select_xai_model(
|
|
"fake-key",
|
|
policy="latest"
|
|
)
|
|
self.assertEqual(result, "grok-4-latest")
|
|
|
|
def test_stable_policy(self):
|
|
# Clear cache first to avoid interference
|
|
from lib import cache
|
|
cache.MODEL_CACHE_FILE.unlink(missing_ok=True)
|
|
result = models.select_xai_model(
|
|
"fake-key",
|
|
policy="stable"
|
|
)
|
|
self.assertEqual(result, "grok-4")
|
|
|
|
def test_pinned_policy(self):
|
|
result = models.select_xai_model(
|
|
"fake-key",
|
|
policy="pinned",
|
|
pin="grok-3"
|
|
)
|
|
self.assertEqual(result, "grok-3")
|
|
|
|
|
|
class TestGetModels(unittest.TestCase):
|
|
def test_no_keys_returns_none(self):
|
|
config = {}
|
|
result = models.get_models(config)
|
|
self.assertIsNone(result["openai"])
|
|
self.assertIsNone(result["xai"])
|
|
|
|
def test_openai_key_only(self):
|
|
config = {"OPENAI_API_KEY": "sk-test"}
|
|
mock_models = [{"id": "gpt-5.2", "created": 1704067200}]
|
|
result = models.get_models(config, mock_openai_models=mock_models)
|
|
self.assertEqual(result["openai"], "gpt-5.2")
|
|
self.assertIsNone(result["xai"])
|
|
|
|
def test_both_keys(self):
|
|
config = {
|
|
"OPENAI_API_KEY": "sk-test",
|
|
"XAI_API_KEY": "xai-test",
|
|
}
|
|
mock_openai = [{"id": "gpt-5.2", "created": 1704067200}]
|
|
mock_xai = [{"id": "grok-4-latest", "created": 1704067200}]
|
|
result = models.get_models(config, mock_openai, mock_xai)
|
|
self.assertEqual(result["openai"], "gpt-5.2")
|
|
self.assertEqual(result["xai"], "grok-4-latest")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|