sashatrask
73 lines
3.0 KiB
Python
73 lines
3.0 KiB
Python
from pathlib import Path
|
|
import sys
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
APP_DIR = ROOT / "app"
|
|
sys.path.insert(0, str(APP_DIR))
|
|
|
|
import keycheck_runner
|
|
from keycheckers.gemini import geminiKeycheck as gemini
|
|
from keycheckers.openai import Keycheck as openai
|
|
from keycheckers.xai import xaiKeycheck as xai
|
|
|
|
|
|
class ModelProbePreferenceTests(unittest.TestCase):
|
|
def test_gemini_prefers_current_pro_then_current_flash(self):
|
|
models = ["gemini-2.5-flash", "gemini-3.7-flash", "gemini-3.1-pro-preview"]
|
|
self.assertEqual(gemini.choose_probe_model(models), "gemini-3.1-pro-preview")
|
|
self.assertEqual(
|
|
gemini.choose_probe_model(["gemini-2.5-flash", "gemini-3.7-flash"]),
|
|
"gemini-3.7-flash",
|
|
)
|
|
self.assertEqual(gemini.choose_probe_model(["gemini-4-preview"]), "gemini-4-preview")
|
|
|
|
def test_xai_prefers_grok_4_6_and_falls_back_to_advertised_model(self):
|
|
self.assertEqual(xai.choose_chat_model(["grok-4.5", "grok-4.6"]), "grok-4.6")
|
|
self.assertEqual(
|
|
xai.choose_chat_model(["grok-4.20-0309-non-reasoning", "grok-4.5"]),
|
|
"grok-4.5",
|
|
)
|
|
|
|
def test_openai_prefers_advertised_gpt_5_6_sol(self):
|
|
models = ["gpt-5", "gpt-5.6-terra", "gpt-5.6-sol", "o3-pro"]
|
|
self.assertEqual(openai.choose_probe_model(models), "gpt-5.6-sol")
|
|
self.assertEqual(openai.choose_probe_model(["gpt-5", "o3-pro"]), "gpt-5")
|
|
self.assertEqual(openai.choose_probe_model(["gpt-4o", "text-embedding-3-large"]), "gpt-4o")
|
|
self.assertEqual(
|
|
openai.reportable_target_models(models),
|
|
{"gpt-5", "gpt-5.6-terra", "gpt-5.6-sol", "o3-pro"},
|
|
)
|
|
|
|
def test_openai_probe_uses_supported_completion_budget(self):
|
|
response = mock.Mock(status_code=200, text='{"choices":[{}]}')
|
|
response.json.return_value = {"choices": [{}], "service_tier": "default"}
|
|
with mock.patch.object(openai.requests, "post", return_value=response) as request:
|
|
status, tier = openai.check_balance_and_tier("fixture-key", "gpt-5.6-sol", None)
|
|
|
|
self.assertEqual((status, tier), ("ok", "default"))
|
|
payload = request.call_args.kwargs["json"]
|
|
self.assertEqual(payload["model"], "gpt-5.6-sol")
|
|
self.assertEqual(payload["max_completion_tokens"], 16)
|
|
self.assertNotIn("max_tokens", payload)
|
|
|
|
def test_postgres_projection_detail_includes_probe_and_inventory(self):
|
|
detail = keycheck_runner.postgres_status_projection_detail({
|
|
"message": "accepted",
|
|
"llm_probe_model": "gpt-5.6-sol",
|
|
"llm_probe_status": "GENERATION_OK",
|
|
"model_count": 3,
|
|
"model_inventory": ["gpt-5.6-sol", "gpt-5", "o3-pro"],
|
|
})
|
|
self.assertEqual(
|
|
detail,
|
|
"accepted; probe_model=gpt-5.6-sol; probe_status=GENERATION_OK; "
|
|
"model_count=3; models=gpt-5.6-sol,gpt-5,o3-pro",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|