import importlib.util import io import json import threading import unittest from pathlib import Path import requests from requests.adapters import BaseAdapter from urllib3.response import HTTPResponse module_path = Path(__file__).parents[1] / "utils" / "api_audit.py" spec = importlib.util.spec_from_file_location("api_audit", module_path) api_audit = importlib.util.module_from_spec(spec) spec.loader.exec_module(api_audit) class ApiAuditTests(unittest.TestCase): def tearDown(self): api_audit.clear_current_job() api_audit._recorder = None def test_records_original_request_and_response_bodies(self): records = [] api_audit._recorder = lambda **record: records.append(record) api_audit.set_current_job("job-1") request = type( "Request", (), { "method": "POST", "url": "https://api.example/tasks?token=secret&mode=video", "headers": {"Authorization": "Bearer secret", "Content-Type": "application/json"}, "body": b'{"prompt":"original"}', }, )() response = type( "Response", (), { "status_code": 200, "headers": {"Content-Type": "application/json"}, "content": b'{"task_id":"task-123","status":"running"}', }, )() api_audit._record("requests", request, response) self.assertEqual(1, len(records)) self.assertEqual("job-1", records[0]["prompt_id"]) self.assertEqual(b'{"prompt":"original"}', records[0]["request_body"]) self.assertEqual( b'{"task_id":"task-123","status":"running"}', records[0]["response_body"], ) self.assertIn("token=%5BREDACTED%5D", records[0]["url"]) self.assertEqual("[REDACTED]", json.loads(records[0]["request_headers"])["Authorization"]) def test_binary_media_response_is_replaced_with_small_reference(self): records = [] api_audit._recorder = lambda **record: records.append(record) api_audit.set_current_job("job-video") video = b"\x00\x00\x00\x18ftypmp42" + (b"x" * 100000) request = type( "Request", (), {"method": "GET", "url": "https://oss.example/output.mp4", "headers": {}, "body": None}, )() response = type( "Response", (), {"status_code": 200, "headers": {"Content-Type": "video/mp4"}, "content": video}, )() api_audit._record("requests", request, response) reference = json.loads(records[0]["response_body"])["$account_manager_body_ref"] self.assertFalse(reference["stored"]) self.assertEqual("video/mp4", reference["content_type"]) self.assertEqual(len(video), reference["content_bytes"]) self.assertEqual(api_audit.hashlib.sha256(video).hexdigest(), reference["sha256"]) self.assertLess(len(records[0]["response_body"]), 300) def test_json_keeps_api_fields_but_removes_embedded_base64_media(self): records = [] api_audit._recorder = lambda **record: records.append(record) api_audit.set_current_job("job-image") image_data = "a" * (api_audit.LARGE_MEDIA_STRING_BYTES + 1) body = json.dumps( {"prompt": "keep me", "first_frame_image": image_data} ).encode() request = type( "Request", (), { "method": "POST", "url": "https://api.example/tasks", "headers": {"Content-Type": "application/json"}, "body": body, }, )() response = type( "Response", (), { "status_code": 200, "headers": {"Content-Type": "application/json"}, "content": b'{"task_id":"task-123","status":"running"}', }, )() api_audit._record("requests", request, response) stored_request = json.loads(records[0]["request_body"]) stored_response = json.loads(records[0]["response_body"]) self.assertEqual("keep me", stored_request["prompt"]) self.assertNotIn(image_data, records[0]["request_body"].decode()) media = stored_request["first_frame_image"]["$account_manager_media_ref"] self.assertEqual(len(image_data), media["encoded_bytes"]) self.assertEqual("task-123", stored_response["task_id"]) self.assertEqual("running", stored_response["status"]) def test_multipart_request_body_is_not_stored(self): raw = b"multipart media" * 10000 stored = api_audit._audit_body( raw, {"Content-Type": "multipart/form-data; boundary=example"} ) reference = json.loads(stored)["$account_manager_body_ref"] self.assertEqual("multipart_body", reference["reason"]) self.assertEqual(len(raw), reference["content_bytes"]) self.assertLess(len(stored), 300) def test_background_thread_uses_single_worker_fallback_job(self): seen = [] api_audit.set_current_job("job-thread") thread = threading.Thread(target=lambda: seen.append(api_audit.current_job())) thread.start() thread.join() self.assertEqual(["job-thread"], seen) def test_installed_hooks_capture_requests_and_httpx(self): import httpx import requests from requests.adapters import BaseAdapter records = [] api_audit.install_api_audit(lambda **record: records.append(record)) api_audit.set_current_job("hook-job") class Adapter(BaseAdapter): def send(self, request, **kwargs): response = requests.Response() response.status_code = 200 response.headers["Content-Type"] = "application/json" response._content = b'{"task_id":"requests-task"}' response.request = request return response def close(self): pass session = requests.Session() session.mount("https://", Adapter()) session.post("https://requests.example/tasks", json={"prompt": "one"}) def handler(request): return httpx.Response(200, json={"task_id": "httpx-task"}) with httpx.Client(transport=httpx.MockTransport(handler)) as client: client.post("https://httpx.example/tasks", json={"prompt": "two"}) self.assertEqual(["requests", "httpx"], [record["client"] for record in records]) self.assertIn(b"requests-task", records[0]["response_body"]) self.assertIn(b"httpx-task", records[1]["response_body"]) def test_streamed_response_remains_available_to_raw_downloaders(self): records = [] api_audit.install_api_audit(lambda **record: records.append(record)) api_audit.set_current_job("cos-download") video = b"\x00\x00\x00\x18ftypmp42" + b"x" * 100000 class Adapter(BaseAdapter): def send(self, request, **kwargs): response = requests.Response() response.status_code = 200 response.headers.update({"Content-Type": "video/mp4", "Content-Length": str(len(video))}) response.raw = HTTPResponse(body=io.BytesIO(video), preload_content=False) response.request = request return response def close(self): pass for session_default in (False, True): with self.subTest(session_default=session_default), requests.Session() as session: session.mount("https://", Adapter()) session.stream = session_default kwargs = {} if session_default else {"stream": True} with session.get("https://cos.example/output.mp4", **kwargs) as response: self.assertFalse(response._content_consumed) self.assertEqual(video, b"".join(response.raw.stream(1024))) self.assertIsNone(records[-1]["response_body"]) self.assertEqual(200, records[-1]["response_status"]) self.assertEqual(str(len(video)), json.loads(records[-1]["response_headers"])["Content-Length"]) if __name__ == "__main__": unittest.main()