mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
test(proxy): exercise tinyfish route through the app with a faked upstream
This commit is contained in:
Zachary Lyon
parent
fec4678013
commit
6692946d70
1 changed files with 64 additions and 112 deletions
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue