Spaces:
Running
Running
| import unittest | |
| from unittest.mock import AsyncMock, patch | |
| from uuid import UUID | |
| import httpx | |
| from fastapi import HTTPException | |
| from starlette.requests import Request | |
| import video_stream | |
| class ProbeStream(httpx.AsyncByteStream): | |
| def __init__(self): | |
| self.reads = 0 | |
| self.closed = False | |
| async def __aiter__(self): | |
| for chunk in (b"first", b"second"): | |
| self.reads += 1 | |
| yield chunk | |
| async def aclose(self): | |
| self.closed = True | |
| def request(method="GET", headers=None): | |
| return Request({"type": "http", "method": method, "path": "/storage/test/videos/camera_1.mp4", | |
| "headers": [(key.encode(), value.encode()) for key, value in (headers or {}).items()]}) | |
| class VideoStreamingTest(unittest.IsolatedAsyncioTestCase): | |
| async def test_streams_ranges_lazily_and_closes_after_client_cancels(self): | |
| body = ProbeStream() | |
| seen = [] | |
| def handler(req): | |
| seen.append(req) | |
| return httpx.Response(206, headers={"Content-Range": "bytes 0-10/1000000", "Content-Length": "11", | |
| "Accept-Ranges": "bytes", "Content-Type": "video/mp4"}, stream=body) | |
| client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) | |
| with patch.object(video_stream.httpx, "AsyncClient", return_value=client): | |
| response = await video_stream.video_response(request(headers={"range": "bytes=0-10", "cookie": "private-cookie"}), | |
| "test/storage", "sessions/id/videos/camera_1.mp4", "storage-token") | |
| self.assertEqual(body.reads, 0) | |
| self.assertEqual(response.status_code, 206) | |
| self.assertEqual(response.headers["content-range"], "bytes 0-10/1000000") | |
| self.assertEqual(seen[0].headers["range"], "bytes=0-10") | |
| self.assertNotIn("cookie", seen[0].headers) | |
| self.assertEqual(await anext(response.body_iterator), b"first") | |
| await response.body_iterator.aclose() | |
| await response.background() | |
| self.assertEqual(body.reads, 1) | |
| self.assertTrue(body.closed) | |
| self.assertTrue(client.is_closed) | |
| async def test_cdn_redirect_strips_authorization_and_keeps_range(self): | |
| seen = [] | |
| def handler(req): | |
| seen.append(req) | |
| if req.url.host == "huggingface.co": | |
| return httpx.Response(302, headers={"Location": "https://cdn.example/video?signature=test"}) | |
| return httpx.Response(206, headers={"Content-Range": "bytes 0-1/10"}, stream=ProbeStream()) | |
| client = httpx.AsyncClient(transport=httpx.MockTransport(handler), follow_redirects=True) | |
| with patch.object(video_stream.httpx, "AsyncClient", return_value=client): | |
| response = await video_stream.video_response(request(headers={"range": "bytes=0-1"}), | |
| "test/storage", "video.mp4", "storage-token") | |
| self.assertEqual(seen[0].headers["authorization"], "Bearer storage-token") | |
| self.assertNotIn("authorization", seen[1].headers) | |
| self.assertEqual(seen[1].headers["range"], "bytes=0-1") | |
| await response.background() | |
| async def test_head_preserves_size_without_reading_video(self): | |
| body = ProbeStream() | |
| client = httpx.AsyncClient(transport=httpx.MockTransport(lambda req: httpx.Response( | |
| 200, headers={"Content-Length": "1000000", "Accept-Ranges": "bytes"}, stream=body))) | |
| with patch.object(video_stream.httpx, "AsyncClient", return_value=client): | |
| response = await video_stream.video_response(request("HEAD"), "test/storage", "video.mp4", "token") | |
| self.assertEqual(response.headers["content-length"], "1000000") | |
| self.assertEqual(response.body, b"") | |
| self.assertEqual(body.reads, 0) | |
| self.assertTrue(client.is_closed) | |
| async def test_invalid_range_keeps_total_size_and_omits_storage_error_body(self): | |
| client = httpx.AsyncClient(transport=httpx.MockTransport(lambda req: httpx.Response( | |
| 416, headers={"Content-Range": "bytes */100", "Content-Length": "100"}, stream=ProbeStream()))) | |
| with patch.object(video_stream.httpx, "AsyncClient", return_value=client): | |
| response = await video_stream.video_response(request(headers={"range": "bytes=999-"}), "test/storage", "video.mp4", "token") | |
| self.assertEqual(response.status_code, 416) | |
| self.assertEqual(response.headers["content-range"], "bytes */100") | |
| self.assertEqual(response.headers["content-length"], "0") | |
| self.assertTrue(client.is_closed) | |
| async def test_upstream_errors_close_connections_and_hide_storage_details(self): | |
| for status, expected in ((404, 404), (401, 502), (503, 502)): | |
| with self.subTest(status=status): | |
| client = httpx.AsyncClient(transport=httpx.MockTransport(lambda req: httpx.Response(status, stream=ProbeStream()))) | |
| with patch.object(video_stream.httpx, "AsyncClient", return_value=client): | |
| with self.assertRaises(HTTPException) as error: | |
| await video_stream.video_response(request(), "test/storage", "video.mp4", "private-token") | |
| self.assertEqual(error.exception.status_code, expected) | |
| self.assertNotIn("private-token", error.exception.detail) | |
| self.assertTrue(client.is_closed) | |
| async def test_video_route_checks_ownership_before_streaming(self): | |
| import app as analysis_app | |
| req = request() | |
| req.state.user_id = "11111111-1111-4111-8111-111111111111" | |
| req.state.is_admin = False | |
| with patch.object(analysis_app, "session_row", side_effect=HTTPException(404, "Session not found")), \ | |
| patch.object(analysis_app, "video_response", new_callable=AsyncMock) as stream: | |
| with self.assertRaises(HTTPException): | |
| await analysis_app.video(UUID("22222222-2222-4222-8222-222222222222"), "camera_1.mp4", req) | |
| stream.assert_not_called() | |
| if __name__ == "__main__": | |
| unittest.main() | |