From 9fafa91d9fc7bb5098d5da333eb28fc9b104e456 Mon Sep 17 00:00:00 2001
From: Deepanshu Pal <40927968+DeepanshuPal@users.noreply.github.com>
Date: Fri, 2 Oct 2026 21:37:31 +0530
Subject: [PATCH] Implement test for custom LLM provider in websocket

Add test for _aresponses_websocket to check custom LLM provider usage.
---
 .../test_responses_api_request_body.py        | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/tests/unit/responses/test_responses_api_request_body.py b/tests/unit/responses/test_responses_api_request_body.py
index b27401d693a..2705fd0ca7e 100644
--- a/tests/unit/responses/test_responses_api_request_body.py
+++ b/tests/unit/responses/test_responses_api_request_body.py
@@ -532,6 +532,29 @@ async def test_aresponses_websocket_strips_responses_routing_prefix_from_openai_
         assert mock_ws.call_args.kwargs["custom_llm_provider"] == "openai"
 
 
+@pytest.mark.asyncio
+async def test_aresponses_websocket_uses_custom_llm_provider_from_kwargs():
+    from unittest.mock import MagicMock
+
+    from litellm.responses.main import _aresponses_websocket
+
+    with patch.object(
+        import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket",
+        new_callable=AsyncMock,
+    ) as mock_ws:
+        await _aresponses_websocket(
+            model="model-x",
+            websocket=MagicMock(),
+            api_key="sk-test",
+            custom_llm_provider="openai",
+            litellm_logging_obj=MagicMock(),
+        )
+
+        mock_ws.assert_awaited_once()
+        assert mock_ws.call_args.kwargs["custom_llm_provider"] == "openai"
+        assert mock_ws.call_args.kwargs["model"] == "model-x"
+
+
 @pytest.mark.asyncio
 async def test_aresponses_websocket_keeps_routing_hints_out_of_the_relay_kwargs():  # test-quality-ok: the relay kwargs are the only place a dropped key is observable; the provider socket behind them is the boundary
     from unittest.mock import MagicMock