diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 9724edf4a56..09985af9660 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -41,7 +41,6 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( milvus_proxy_route, mistral_proxy_route, openai_proxy_route, - tinyfish_proxy_route, vertex_discovery_proxy_route, vertex_proxy_route, vllm_proxy_route, @@ -5496,139 +5495,92 @@ class TestAzureRelayDeploymentSegment: class TestTinyFishProxyRoute: - """Tests for the TinyFish Agent pass-through route.""" + """Tests for the TinyFish Agent pass-through route, faking the upstream HTTP boundary.""" - def _mock_request(self, method: str, body: bytes = b"") -> MagicMock: - mock_request = MagicMock(spec=Request) - mock_request.method = method - mock_request.query_params = {} - mock_request.headers = {} - mock_request.body = AsyncMock(return_value=body) - return mock_request + RUN_BODY = {"url": "https://scrapeme.live/shop", "goal": "Extract the first 2 product names. Return JSON."} - @pytest.mark.asyncio - async def test_forwards_allowed_run_endpoint_with_server_key(self): - with ( - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", - return_value="sk-tf-server", - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" - ) as mock_create_route, - ): - mock_endpoint_func = AsyncMock(return_value={"run_id": "run-1", "status": "COMPLETED"}) - mock_create_route.return_value = mock_endpoint_func + @pytest.fixture + def tinyfish_client(self, monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]: + from litellm.proxy.proxy_server import app - result = await tinyfish_proxy_route( - endpoint="v1/automation/run", - request=self._mock_request("POST", b'{"url": "https://scrapeme.live/shop", "goal": "extract"}'), - fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(), + monkeypatch.setenv("TINYFISH_API_KEY", "sk-tf-upstream") + monkeypatch.delenv("TINYFISH_AGENT_API_BASE", raising=False) + monkeypatch.delenv("TINYFISH_ALLOW_AUTHENTICATED_RUNS", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual")) + yield TestClient(app) + + def test_forwards_run_with_server_key_not_callers(self, tinyfish_client: TestClient) -> None: + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post("https://agent.tinyfish.ai/v1/automation/run").mock( + return_value=httpx.Response(200, json={"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 2}) + ) + response = tinyfish_client.post( + "/tinyfish/v1/automation/run", json=self.RUN_BODY, headers={"X-API-Key": "sk-callers-virtual-key"} ) - call_args = mock_create_route.call_args[1] - assert call_args["target"] == "https://agent.tinyfish.ai/v1/automation/run" - assert dict(call_args["custom_headers"]) == {"X-API-Key": "sk-tf-server"} - assert call_args["custom_llm_provider"] == "tinyfish" - assert result == {"run_id": "run-1", "status": "COMPLETED"} + assert (response.status_code, response.json()["run_id"]) == (200, "run-1") + assert route.calls.last.request.headers["x-api-key"] == "sk-tf-upstream" - @pytest.mark.asyncio @pytest.mark.parametrize( - "method,endpoint", + "method,path", [ - ("GET", "v1/vault/items"), - ("GET", "v1/wallet"), - ("POST", "v1/browser-profiles"), - ("GET", "v1/automation/run"), + ("GET", "/tinyfish/v1/vault/items"), + ("GET", "/tinyfish/v1/wallet"), + ("POST", "/tinyfish/v1/browser-profiles"), + ("GET", "/tinyfish/v1/automation/run"), ], ) - async def test_blocks_endpoints_outside_allowlist(self, method, endpoint): - with pytest.raises(HTTPException) as exc_info: - await tinyfish_proxy_route( - endpoint=endpoint, - request=self._mock_request(method), - fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(), - ) + def test_blocks_endpoints_outside_allowlist(self, tinyfish_client: TestClient, method: str, path: str) -> None: + with respx.mock: + response = tinyfish_client.request(method, path) - assert exc_info.value.status_code == 403 + assert response.status_code == 403 + assert "not an allowed TinyFish Agent passthrough endpoint" in response.json()["detail"] - @pytest.mark.asyncio - async def test_rejects_authenticated_run_fields_by_default(self, monkeypatch): - monkeypatch.delenv("TINYFISH_ALLOW_AUTHENTICATED_RUNS", raising=False) + def test_rejects_authenticated_run_fields_by_default(self, tinyfish_client: TestClient) -> None: + with respx.mock: + response = tinyfish_client.post("/tinyfish/v1/automation/run", json={**self.RUN_BODY, "use_vault": True}) - with pytest.raises(HTTPException) as exc_info: - await tinyfish_proxy_route( - endpoint="v1/automation/run", - request=self._mock_request("POST", b'{"url": "https://x.com", "goal": "g", "use_vault": true}'), - fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(), - ) + assert response.status_code == 403 + assert "use_vault" in response.json()["detail"] - assert exc_info.value.status_code == 403 - assert "use_vault" in exc_info.value.detail - - @pytest.mark.asyncio - async def test_env_opt_in_allows_authenticated_run_fields(self, monkeypatch): + def test_env_opt_in_allows_authenticated_run_fields( + self, tinyfish_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: monkeypatch.setenv("TINYFISH_ALLOW_AUTHENTICATED_RUNS", "true") - with ( - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", - return_value="sk-tf-server", - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" - ) as mock_create_route, - ): - mock_create_route.return_value = AsyncMock(return_value={"ok": True}) - - result = await tinyfish_proxy_route( - endpoint="v1/automation/run", - request=self._mock_request("POST", b'{"url": "https://x.com", "goal": "g", "use_vault": true}'), - fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(), + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post("https://agent.tinyfish.ai/v1/automation/run").mock( + return_value=httpx.Response(200, json={"run_id": "run-2", "status": "COMPLETED", "num_of_steps": 1}) ) + response = tinyfish_client.post("/tinyfish/v1/automation/run", json={**self.RUN_BODY, "use_vault": True}) - assert result == {"ok": True} + assert response.status_code == 200 + assert json.loads(route.calls.last.request.content)["use_vault"] is True - @pytest.mark.asyncio - async def test_raises_401_on_missing_api_key(self): - with patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", - return_value=None, - ): - with pytest.raises(HTTPException) as exc_info: - await tinyfish_proxy_route( - endpoint="v1/runs", - request=self._mock_request("GET"), - fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(), - ) + def test_returns_401_on_missing_api_key( + self, tinyfish_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("TINYFISH_API_KEY") - assert exc_info.value.status_code == 401 + with respx.mock: + response = tinyfish_client.get("/tinyfish/v1/runs") - @pytest.mark.asyncio - async def test_env_base_override_changes_target(self, monkeypatch): + assert response.status_code == 401 + assert "TINYFISH_API_KEY" in response.json()["detail"] + + def test_env_base_override_changes_target( + self, tinyfish_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: monkeypatch.setenv("TINYFISH_AGENT_API_BASE", "https://agent.staging.tinyfish.ai") - with ( - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", - return_value="sk-tf-server", - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" - ) as mock_create_route, - ): - mock_create_route.return_value = AsyncMock(return_value={}) - - await tinyfish_proxy_route( - endpoint="v1/runs/run-123", - request=self._mock_request("GET"), - fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(), + with respx.mock(assert_all_called=True) as upstream: + upstream.get("https://agent.staging.tinyfish.ai/v1/runs/run-123").mock( + return_value=httpx.Response(200, json={"run_id": "run-123", "status": "RUNNING"}) ) + response = tinyfish_client.get("/tinyfish/v1/runs/run-123") - assert mock_create_route.call_args[1]["target"] == "https://agent.staging.tinyfish.ai/v1/runs/run-123" + assert (response.status_code, response.json()["status"]) == (200, "RUNNING")