import pytest from unittest.mock import AsyncMock, patch from engine import run_comparison, pick_best_value from config import MODELS def test_pick_best_value_cheapest_above_threshold(): scores = {"nano": 8, "mini": 9, "large": 9, "reasoning": "all good"} costs = {"nano": 0.0001, "mini": 0.0003, "large": 0.0} best = pick_best_value(scores, costs, threshold=7) assert best == "large" def test_pick_best_value_free_model_scores_below_threshold(): scores = {"nano": 5, "mini": 8, "large": 6, "reasoning": "nano and large weak"} costs = {"nano": 0.0001, "mini": 0.0003, "large": 0.0} best = pick_best_value(scores, costs, threshold=7) assert best == "mini" def test_pick_best_value_none_above_threshold(): scores = {"nano": 3, "mini": 5, "large": 6, "reasoning": "all weak"} costs = {"nano": 0.0001, "mini": 0.0003, "large": 0.0} best = pick_best_value(scores, costs, threshold=7) assert best is None def test_pick_best_value_tie_goes_to_cheapest(): scores = {"nano": 8, "mini": 8, "large": 8, "reasoning": "all equal"} costs = {"nano": 0.0001, "mini": 0.0003, "large": 0.0} best = pick_best_value(scores, costs, threshold=7) assert best == "large" @pytest.mark.asyncio async def test_run_comparison_prompt_too_long(): result = await run_comparison("x" * 2001, ip="1.2.3.4") assert result["error"] is not None assert "2,000" in result["error"] @pytest.mark.asyncio async def test_run_comparison_success(): mock_model_results = { "nano": {"content": "Nano response", "prompt_tokens": 10, "completion_tokens": 20, "error": None}, "mini": {"content": "Mini response", "prompt_tokens": 10, "completion_tokens": 20, "error": None}, "large": {"content": "Large response", "prompt_tokens": 10, "completion_tokens": 20, "error": None}, } mock_scores = {"nano": 8, "mini": 9, "large": 9, "reasoning": "all good"} with patch("engine.call_models_parallel", new_callable=AsyncMock, return_value=mock_model_results), \ patch("engine.judge_responses", new_callable=AsyncMock, return_value=mock_scores), \ patch("engine.rate_limiter") as mock_limiter: mock_limiter.check.return_value = True result = await run_comparison("Write hello", ip="1.2.3.4") assert result["error"] is None assert result["responses"]["nano"]["content"] == "Nano response" assert result["scores"]["nano"] == 8 assert result["best_value"] is not None @pytest.mark.asyncio async def test_run_comparison_all_models_fail(): mock_model_results = { "nano": {"content": None, "prompt_tokens": 0, "completion_tokens": 0, "error": "timeout"}, "mini": {"content": None, "prompt_tokens": 0, "completion_tokens": 0, "error": "timeout"}, "large": {"content": None, "prompt_tokens": 0, "completion_tokens": 0, "error": "timeout"}, } with patch("engine.call_models_parallel", new_callable=AsyncMock, return_value=mock_model_results), \ patch("engine.rate_limiter") as mock_limiter: mock_limiter.check.return_value = True result = await run_comparison("Test", ip="1.2.3.4") assert result["error"] is not None assert "All models failed" in result["error"]