diff --git a/tests/pass_through_unit_tests/test_custom_logger_passthrough.py b/tests/pass_through_unit_tests/test_custom_logger_passthrough.py index d7add68753b..a7d6378fa18 100644 --- a/tests/pass_through_unit_tests/test_custom_logger_passthrough.py +++ b/tests/pass_through_unit_tests/test_custom_logger_passthrough.py @@ -1,5 +1,6 @@ import json -import os +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Optional from fastapi import Request import pytest @@ -29,14 +30,59 @@ class TestCustomLogger(CustomLogger): self.logged_kwargs = kwargs +UPSTREAM_RESPONSE_BODY = { + "id": "modr-abc123", + "model": "omni-moderation-latest", + "results": [ + { + "flagged": False, + "categories": {"violence": False}, + "category_scores": {"violence": 1.2e-06}, + } + ], +} + + +@pytest.fixture +def upstream(): + received: dict = {} + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self): + body = self.rfile.read(int(self.headers.get("content-length", 0) or 0)) + received["path"] = self.path + received["body"] = json.loads(body or b"{}") + payload = json.dumps(UPSTREAM_RESPONSE_BODY).encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + yield f"http://127.0.0.1:{server.server_port}", received + finally: + server.shutdown() + server.server_close() + + @pytest.mark.asyncio -async def test_untracked_openai_route_passthrough_logging(): +async def test_untracked_openai_route_passthrough_logging(upstream): """Keep this on a route `_is_supported_openai_endpoint` does not claim, or the OpenAI-specific handler takes over and the generic payload stops being exercised.""" + base_url, upstream_received = upstream + test_custom_logger = TestCustomLogger() litellm._async_success_callback = [test_custom_logger] - TARGET_URL = "https://api.openai.com/v1/moderations" + TARGET_URL = f"{base_url}/v1/moderations" REQUEST_BODY = { "model": "omni-moderation-latest", "input": "I want to bake a cake for my friend's birthday.", @@ -52,17 +98,14 @@ async def test_untracked_openai_route_passthrough_logging(): "query_string": b"", "headers": [ (b"content-type", b"application/json"), - ( - b"authorization", - f"Bearer {os.getenv('OPENAI_API_KEY')}".encode(), - ), + (b"authorization", b"Bearer sk-test-passthrough"), ], }, ), target=TARGET_URL, custom_headers={ "Content-Type": "application/json", - "Authorization": f"Bearer {os.getenv('OPENAI_API_KEY')}", + "Authorization": "Bearer sk-test-passthrough", }, user_api_key_dict=UserAPIKeyAuth( api_key="test", @@ -79,6 +122,8 @@ async def test_untracked_openai_route_passthrough_logging(): print("result status code", result.status_code) print("result content", result.body) + assert upstream_received.get("path") == "/v1/moderations" + assert upstream_received.get("body") == REQUEST_BODY assert result.status_code == 200 await asyncio.sleep(1) @@ -93,5 +138,5 @@ async def test_untracked_openai_route_passthrough_logging(): assert passthrough_logging_payload["request_method"] == TARGET_METHOD client_facing_response_body = json.loads(result.body) - assert client_facing_response_body["results"] + assert client_facing_response_body == UPSTREAM_RESPONSE_BODY assert passthrough_logging_payload["response_body"] == client_facing_response_body