Spaces:
Running
Running
File size: 6,145 Bytes
a66365f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | 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()
|