mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-08-09 02:54:37 +00:00
fix: harden azure routing and model override mapping
This commit is contained in:
@@ -1,14 +1,18 @@
|
||||
"""Tests for the model prioritization algorithm."""
|
||||
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
# Set required env vars before importing
|
||||
os.environ["UPSTREAM_BASE_URL"] = "http://test"
|
||||
os.environ["UPSTREAM_API_KEY"] = "test"
|
||||
|
||||
from routstr.algorithm import ( # noqa: E402
|
||||
calculate_model_cost_score,
|
||||
create_model_mappings,
|
||||
get_provider_penalty,
|
||||
)
|
||||
from routstr.payment.models import Architecture, Model, Pricing # noqa: E402
|
||||
@@ -45,11 +49,21 @@ def create_test_model(
|
||||
)
|
||||
|
||||
|
||||
def create_test_provider(name: str, base_url: str = "http://test.com") -> Mock:
|
||||
def create_test_provider(
|
||||
name: str,
|
||||
base_url: str = "http://test.com",
|
||||
*,
|
||||
db_id: int | None = None,
|
||||
models: list[Model] | None = None,
|
||||
upstream_name: str | None = None,
|
||||
) -> Mock:
|
||||
"""Helper to create a test provider mock."""
|
||||
provider = Mock()
|
||||
provider.provider_type = name
|
||||
provider.base_url = base_url
|
||||
provider.db_id = db_id
|
||||
provider.upstream_name = upstream_name or name
|
||||
provider.get_cached_models.return_value = models or []
|
||||
return provider
|
||||
|
||||
|
||||
@@ -99,3 +113,80 @@ def test_get_provider_penalty_openrouter() -> None:
|
||||
provider = create_test_provider("openrouter", "https://openrouter.ai/api/v1")
|
||||
penalty = get_provider_penalty(provider)
|
||||
assert penalty == 1.001
|
||||
|
||||
|
||||
def test_create_model_mappings_includes_db_override_for_missing_cached_model(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Model overrides should still map when provider discovery misses the model."""
|
||||
provider = create_test_provider(
|
||||
"azure",
|
||||
"https://example.openai.azure.com/openai/v1",
|
||||
db_id=7,
|
||||
models=[],
|
||||
)
|
||||
override_model = create_test_model("azure/gpt-4o")
|
||||
override_model.canonical_slug = "azure-deployment"
|
||||
|
||||
def fake_row_to_model(*args, **kwargs) -> Model: # type: ignore[no-untyped-def]
|
||||
return override_model
|
||||
|
||||
monkeypatch.setattr("routstr.payment.models._row_to_model", fake_row_to_model)
|
||||
|
||||
override_row = SimpleNamespace(id="azure/gpt-4o", upstream_provider_id=7, enabled=True)
|
||||
|
||||
model_instances, provider_map, unique_models = create_model_mappings(
|
||||
upstreams=[provider],
|
||||
overrides_by_id={"azure/gpt-4o": (override_row, 1.01)},
|
||||
disabled_model_ids=set(),
|
||||
)
|
||||
|
||||
assert "azure/gpt-4o" in model_instances
|
||||
assert provider_map["azure/gpt-4o"] == [provider]
|
||||
assert "gpt-4o" in unique_models
|
||||
|
||||
|
||||
def test_create_model_mappings_dedupes_with_provider_identity_not_provider_type(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Different provider instances of same type should both survive dedupe."""
|
||||
provider_a_model = create_test_model(
|
||||
"azure/gpt-4o", prompt_price=0.01, completion_price=0.01
|
||||
)
|
||||
provider_a = create_test_provider(
|
||||
"azure",
|
||||
"https://a.openai.azure.com/openai/v1",
|
||||
db_id=1,
|
||||
models=[provider_a_model],
|
||||
upstream_name="azure-a",
|
||||
)
|
||||
provider_b = create_test_provider(
|
||||
"azure",
|
||||
"https://b.openai.azure.com/openai/v1",
|
||||
db_id=2,
|
||||
models=[],
|
||||
upstream_name="azure-b",
|
||||
)
|
||||
|
||||
override_model = create_test_model(
|
||||
"azure/gpt-4o", prompt_price=0.001, completion_price=0.001
|
||||
)
|
||||
override_model.canonical_slug = "azure-b-deployment"
|
||||
|
||||
def fake_row_to_model(*args, **kwargs) -> Model: # type: ignore[no-untyped-def]
|
||||
return override_model
|
||||
|
||||
monkeypatch.setattr("routstr.payment.models._row_to_model", fake_row_to_model)
|
||||
|
||||
override_row = SimpleNamespace(id="azure/gpt-4o", upstream_provider_id=2, enabled=True)
|
||||
|
||||
_, provider_map, _ = create_model_mappings(
|
||||
upstreams=[provider_a, provider_b],
|
||||
overrides_by_id={"azure/gpt-4o": (override_row, 1.01)},
|
||||
disabled_model_ids=set(),
|
||||
)
|
||||
|
||||
providers_for_alias = provider_map["azure/gpt-4o"]
|
||||
assert provider_a in providers_for_alias
|
||||
assert provider_b in providers_for_alias
|
||||
assert len(providers_for_alias) == 2
|
||||
|
||||
Reference in New Issue
Block a user