From 6809e8087dbd3e656c15c9931d88ebd627fd96ea Mon Sep 17 00:00:00 2001
From: kerry <kerry@berri.ai>
Date: Tue, 22 Sep 2026 21:21:50 +0000
Subject: [PATCH] test(integration): keep the selected model in the response
 and spend log for an Azure Model Router alias (Pylon #7636)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 tests/integration/contracts.json              |  3 +
 .../spend/test_model_router_selected_model.py | 74 +++++++++++++++++++
 2 files changed, 77 insertions(+)
 create mode 100644 tests/integration/spend/test_model_router_selected_model.py

diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json
index d12e3ae4620..0b7321a7924 100644
--- a/tests/integration/contracts.json
+++ b/tests/integration/contracts.json
@@ -272,6 +272,9 @@
     "tests/integration/spend/test_filtered_ledger.py::test_rotated_keys_users_and_model_groups_preserve_success_failure_cache_ledger": [
       "quota_management.spend_tracking.filtered_ledger_preserves_owner_identity_and_totals"
     ],
+    "tests/integration/spend/test_model_router_selected_model.py::test_model_router_alias_without_router_in_name_keeps_selected_model_in_response_and_spend_log": [
+      "spend.model_router.selected_model_is_returned_and_persisted_for_plain_alias"
+    ],
     "tests/integration/management/test_partial_update_sequences.py::test_restricted_actor_cannot_detach_key_from_project": [
       "mgmt.key.update.project_detach_denied_to_restricted_actor"
     ],
diff --git a/tests/integration/spend/test_model_router_selected_model.py b/tests/integration/spend/test_model_router_selected_model.py
new file mode 100644
index 00000000000..a95ee4f9feb
--- /dev/null
+++ b/tests/integration/spend/test_model_router_selected_model.py
@@ -0,0 +1,74 @@
+import json
+import uuid
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway, eventually
+from integration._support.database import read_rows
+from integration._support.wire import Reply, Request, wire_server
+
+ROUTER_DEPLOYMENT: Final = "router-deploy"
+SELECTED_MODEL: Final = "grok-4-1-fast-reasoning"
+SELECTED_MODEL_WITH_PROVIDER: Final = f"azure_ai/{SELECTED_MODEL}"
+
+
+@pytest.mark.covers("spend.model_router.selected_model_is_returned_and_persisted_for_plain_alias")
+def test_model_router_alias_without_router_in_name_keeps_selected_model_in_response_and_spend_log(
+    gateway: Gateway,
+) -> None:
+    prompt: Final = uuid.uuid4().hex
+
+    def provider(request: Request) -> Reply:
+        assert request.method == "POST" and request.target == "/chat/completions", request.target
+        assert json.loads(request.body) == {
+            "model": ROUTER_DEPLOYMENT,
+            "messages": [{"role": "user", "content": prompt}],
+            "stream": False,
+        }, request.body
+        return Reply(
+            body=json.dumps(
+                {
+                    "id": "chatcmpl-" + uuid.uuid4().hex,
+                    "object": "chat.completion",
+                    "created": 1,
+                    "model": SELECTED_MODEL,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "message": {"role": "assistant", "content": "routed answer"},
+                            "finish_reason": "stop",
+                        }
+                    ],
+                    "usage": {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40},
+                }
+            ).encode()
+        )
+
+    with wire_server(provider) as wire, gateway.scenario() as scenario:
+        alias: Final = scenario.model(
+            model=f"azure_ai/model_router/{ROUTER_DEPLOYMENT}", api_base=wire.url, num_retries=0
+        )
+        response: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {"model": alias, "messages": [{"role": "user", "content": prompt}]},
+        )
+        assert response.status_code == 200, response.text
+        body: Final = response.json()
+        assert body["model"] == SELECTED_MODEL_WITH_PROVIDER, response.text
+        assert body["choices"][0]["message"]["content"] == "routed answer", response.text
+        assert len(wire.drain()) == 1
+        rows: Final = eventually(
+            lambda: read_rows(
+                'SELECT model, model_group, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
+                (body["id"],),
+            ),
+            lambda values: len(values) == 1,
+            seconds=70,
+        )
+        assert rows == [{"model": SELECTED_MODEL_WITH_PROVIDER, "model_group": alias, "status": "success"}]
+        logs: Final = gateway.request("GET", "/spend/logs", params={"request_id": body["id"]})
+        assert logs.status_code == 200, logs.text
+        assert [(row["model"], row["model_group"]) for row in logs.json()] == [(SELECTED_MODEL_WITH_PROVIDER, alias)], (
+            logs.text
+        )