diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 71af3cf76a6..2b94fdffea9 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -191,15 +191,32 @@ def test_get_model_info_strips_openai_finetune_ids_without_a_custom_suffix(local ], ) def test_get_model_info_falls_back_from_dated_snapshot_to_undated_entry( - local_model_cost_map, model, custom_llm_provider, expected_key -): - info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, + model: str, + custom_llm_provider: str, + expected_key: str, +) -> None: + monkeypatch.delitem(litellm.model_cost, model, raising=False) + monkeypatch.delitem(litellm.model_cost, f"{custom_llm_provider}/{model}", raising=False) + assert expected_key in litellm.model_cost + info: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) assert info["key"] == expected_key -def test_get_model_info_prefers_exact_dated_key_over_stripped(local_model_cost_map): - info = litellm.get_model_info(model="gpt-4o-2024-08-06", custom_llm_provider="openai") - assert info["key"] == "gpt-4o-2024-08-06" +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected_key"), + [ + ("gpt-4o-2024-08-06", "openai", "gpt-4o-2024-08-06"), + ("gpt-5.6-luna-2026-07-09", "azure", "azure/gpt-5.6-luna-2026-07-09"), + ], +) +def test_get_model_info_prefers_exact_dated_key_over_stripped( + local_model_cost_map: None, model: str, custom_llm_provider: str, expected_key: str +) -> None: + assert expected_key in litellm.model_cost + info: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + assert info["key"] == expected_key def test_check_provider_match_azure_ai_allows_openai_and_azure():