| from __future__ import annotations |
|
|
| import asyncio |
| import copy |
| import gzip |
| import json |
| import tempfile |
| import threading |
| import time |
| import unittest |
| from collections import OrderedDict |
| from pathlib import Path |
| from types import SimpleNamespace |
| from unittest.mock import AsyncMock, Mock, patch |
|
|
| IMPORT_ERROR = None |
| try: |
| import aiohttp |
| import app |
| import upstream_http |
| except ModuleNotFoundError as exc: |
| app = None |
| IMPORT_ERROR = exc |
|
|
|
|
| @unittest.skipIf(app is None, f"proxy dependencies unavailable: {IMPORT_ERROR}") |
| class FakeResponse: |
| def __init__(self, status=200, headers=None, chunks=(b"ok",)): |
| self.status = status |
| self.headers = headers or {} |
| self.content = self |
| self.chunks = chunks |
| self.released = False |
|
|
| def release(self): |
| self.released = True |
|
|
| async def read(self, amount=-1): |
| return b"".join(self.chunks)[:amount if amount >= 0 else None] |
|
|
| def iter_chunked(self, size): |
| async def generator(): |
| for chunk in self.chunks: |
| yield chunk |
| return generator() |
|
|
| async def __aenter__(self): |
| return self |
|
|
| async def __aexit__(self, *args): |
| self.release() |
|
|
|
|
| @unittest.skipIf(app is None, f"proxy dependencies unavailable: {IMPORT_ERROR}") |
| class DownloadConnectionWarmTests(unittest.IsolatedAsyncioTestCase): |
| async def asyncSetUp(self): |
| self.response = FakeResponse() |
| self.state = SimpleNamespace( |
| download_warm_enabled=True, download_warm_task=None, |
| download_warm_at=float("-inf"), upstream_semaphore=asyncio.Semaphore(2), |
| http_session=SimpleNamespace(head=Mock(return_value=self.response)), |
| ) |
| self.application = SimpleNamespace(state=self.state) |
|
|
| async def asyncTearDown(self): |
| if self.state.download_warm_task is not None: |
| self.state.download_warm_task.cancel() |
| await asyncio.gather(self.state.download_warm_task, return_exceptions=True) |
|
|
| async def test_ping_returns_before_warmup_and_coalesces_repeated_visits(self): |
| with patch.object(app, "app", self.application): |
| first = await app.api_ping() |
| task = self.state.download_warm_task |
| second = await app.api_ping() |
| self.assertEqual([first.status_code, second.status_code], [204, 204]) |
| self.assertEqual(first.headers["cache-control"], "no-store") |
| self.state.http_session.head.assert_not_called() |
| self.assertIs(self.state.download_warm_task, task) |
| await task |
| self.state.http_session.head.assert_called_once() |
| args, kwargs = self.state.http_session.head.call_args |
| self.assertEqual(args, ("https://huggingface.co/robots.txt",)) |
| self.assertFalse(kwargs["allow_redirects"]) |
| self.assertEqual(kwargs["timeout"].total, 4) |
| self.assertTrue(self.response.released) |
| app.schedule_download_connection_warm(self.application) |
| self.assertIs(self.state.download_warm_task, task) |
| self.assertEqual(self.state.upstream_semaphore._value, 2) |
|
|
| async def test_failed_warmup_backs_off_then_allows_retry(self): |
| self.state.http_session.head.side_effect = aiohttp.ClientConnectionError() |
| app.schedule_download_connection_warm(self.application) |
| await self.state.download_warm_task |
| app.schedule_download_connection_warm(self.application) |
| self.assertEqual(self.state.http_session.head.call_count, 1) |
| self.state.download_warm_at -= app.DOWNLOAD_CONNECTION_WARM_INTERVAL_SECONDS |
| self.state.http_session.head.side_effect = None |
| app.schedule_download_connection_warm(self.application) |
| await self.state.download_warm_task |
| self.assertEqual(self.state.http_session.head.call_count, 2) |
| self.assertEqual(self.state.upstream_semaphore._value, 2) |
|
|
| async def test_saturated_pool_skips_warmup_without_queuing(self): |
| self.state.upstream_semaphore = asyncio.Semaphore(0) |
| await app.warm_download_connection(self.application) |
| self.state.http_session.head.assert_not_called() |
|
|
| async def test_cancellation_releases_pool_and_disabled_state_stops_scheduling(self): |
| entered = asyncio.Event() |
| class StalledResponse(FakeResponse): |
| async def __aenter__(self): |
| entered.set() |
| await asyncio.Future() |
| self.state.http_session.head.return_value = StalledResponse() |
| app.schedule_download_connection_warm(self.application) |
| await entered.wait() |
| task = self.state.download_warm_task |
| task.cancel() |
| with self.assertRaises(asyncio.CancelledError): |
| await task |
| self.assertEqual(self.state.upstream_semaphore._value, 2) |
| self.state.download_warm_enabled = False |
| self.state.download_warm_at = float("-inf") |
| app.schedule_download_connection_warm(self.application) |
| self.assertIs(self.state.download_warm_task, task) |
|
|
|
|
| class FakeSession: |
| def __init__(self, responses): |
| self.responses = list(responses) |
| self.urls = [] |
|
|
| async def get(self, url, **kwargs): |
| self.urls.append((url, kwargs)) |
| response = self.responses.pop(0) |
| if isinstance(response, Exception): |
| raise response |
| return response |
|
|
| async def request(self, method, url, **kwargs): |
| kwargs["method"] = method |
| return await self.get(url, **kwargs) |
|
|
|
|
| @unittest.skipIf(app is None, f"proxy dependencies unavailable: {IMPORT_ERROR}") |
| class AsyncProxyTests(unittest.IsolatedAsyncioTestCase): |
| async def asyncSetUp(self): |
| app.initialize_reader_bucket_state(app.app) |
| self._redirect_cache = dict(upstream_http.upstream_redirect_cache) |
| upstream_http.upstream_redirect_cache.clear() |
|
|
| async def asyncTearDown(self): |
| await app.shutdown_reader_bucket_state(app.app) |
| upstream_http.upstream_redirect_cache.clear() |
| upstream_http.upstream_redirect_cache.update(self._redirect_cache) |
|
|
| def test_reader_assets_decoder_filters_invalid_entries(self): |
| valid_path = "objects/aa/" + "b" * 64 + "/document.pdf" |
| valid_linearized_path = "objects/aa/" + "f" * 64 + "/linearized.pdf" |
| valid_azw_path = "objects/aa/" + "e" * 64 + "/document.azw" |
| page_manifest = "objects/aa/" + "d" * 64 + "/page-manifest.json" |
| media_path = "objects/aa/" + "c" * 64 + "/ffmpeg-audio-mp3-v1/audio.mp3" |
| raw = gzip.compress(json.dumps({"v": 1, "f": { |
| "VoiceOfML/Test\0book.doc": {"s": 2, "m": "p", "p": valid_path}, |
| "VoiceOfML/Test\0book-linearized.pdf": {"s": 2, "m": "p", "p": valid_linearized_path}, |
| "VoiceOfML/Test\0book.azw": {"s": 2, "m": "e", "p": valid_azw_path}, |
| "VoiceOfML/Test\0sound.wma": {"s": 2, "m": "a", "p": media_path}, |
| "VoiceOfML/Test\0scan.pdf": {"s": 2, "m": "p", "p": page_manifest}, |
| "bad-path": {"s": 2, "m": "p", "p": "objects/../../secret.pdf"}, |
| "failed": {"s": 4}, |
| }}).encode()) |
| self.assertEqual(app.decode_reader_assets(raw), {"v": 1, "f": { |
| "VoiceOfML/Test\0book.doc": {"s": 2, "m": "p", "p": valid_path}, |
| "VoiceOfML/Test\0book-linearized.pdf": {"s": 2, "m": "p", "p": valid_linearized_path}, |
| "VoiceOfML/Test\0book.azw": {"s": 2, "m": "e", "p": valid_azw_path}, |
| "VoiceOfML/Test\0sound.wma": {"s": 2, "m": "a", "p": media_path}, |
| "VoiceOfML/Test\0scan.pdf": {"s": 2, "m": "p", "p": page_manifest}, |
| "failed": {"s": 4}, |
| }}) |
|
|
| async def test_reader_assets_refresh_is_cached_and_releases_response(self): |
| payload = gzip.compress(json.dumps({"v": 1, "f": {"failed": {"s": 4}}}).encode()) |
| response = FakeResponse(200, chunks=(payload,)) |
| session = FakeSession([response]) |
| app.app.state.http_session = session |
| app.app.state.reader_assets_lock = asyncio.Lock() |
| old_cache, old_at = app.reader_assets_cache, app.reader_assets_cache_at |
| try: |
| app.reader_assets_cache = {"v": 1, "f": {}} |
| app.reader_assets_cache_at = 0 |
| first = await app.get_reader_assets() |
| second = await app.get_reader_assets() |
| self.assertEqual(first, {"v": 1, "f": {"failed": {"s": 4}}}) |
| self.assertIs(first, second) |
| self.assertEqual(len(session.urls), 1) |
| self.assertTrue(response.released) |
| finally: |
| app.reader_assets_cache, app.reader_assets_cache_at = old_cache, old_at |
|
|
| async def test_search_returns_while_shared_reader_assets_refresh_is_stalled(self): |
| started, release = asyncio.Event(), asyncio.Event() |
| payload = gzip.compress(json.dumps({"v": 1, "f": {"new": {"s": 4}}}).encode()) |
| response = FakeResponse(200, chunks=(payload,)) |
| async def get(*args, **kwargs): |
| started.set() |
| await release.wait() |
| return response |
| app.app.state.http_session = SimpleNamespace(get=AsyncMock(side_effect=get)) |
| app.app.state.reader_assets_lock = asyncio.Lock() |
| app.app.state.reader_assets_task = None |
| with patch.object(app, "reader_assets_cache", {"v": 1, "f": {}}), \ |
| patch.object(app, "reader_assets_cache_at", 0), \ |
| patch.object(app, "search", return_value={"results": [], "total": 0, "page": 1, "page_size": 100}): |
| try: |
| first = await asyncio.wait_for(app.api_search(app.SearchRequest()), 1) |
| await started.wait() |
| task = app.app.state.reader_assets_task |
| second = await asyncio.wait_for(app.api_search(app.SearchRequest()), 1) |
| self.assertEqual([first.status_code, second.status_code], [200, 200]) |
| self.assertIs(task, app.app.state.reader_assets_task) |
| self.assertFalse(task.done()) |
| release.set() |
| await task |
| self.assertEqual(app.reader_assets_cache["f"], {"new": {"s": 4}}) |
| self.assertTrue(response.released) |
| self.assertEqual(app.app.state.http_session.get.call_count, 1) |
| finally: |
| release.set() |
| if app.app.state.reader_assets_task: |
| await app.app.state.reader_assets_task |
|
|
| async def test_reader_assets_refresh_failure_backs_off_with_existing_cache(self): |
| response = FakeResponse(503) |
| session = FakeSession([response]) |
| app.app.state.http_session = session |
| app.app.state.reader_assets_lock = asyncio.Lock() |
| old_cache, old_at = app.reader_assets_cache, app.reader_assets_cache_at |
| try: |
| app.reader_assets_cache = {"v": 1, "f": {"old": {"s": 4}}} |
| app.reader_assets_cache_at = time.monotonic() - app.READER_ASSETS_CACHE_TTL_SECONDS - 1 |
| first = await app.get_reader_assets() |
| second = await app.get_reader_assets() |
| self.assertIs(first, second) |
| self.assertEqual(len(session.urls), 1) |
| self.assertTrue(response.released) |
| finally: |
| app.reader_assets_cache, app.reader_assets_cache_at = old_cache, old_at |
|
|
| async def test_initial_reader_assets_failure_is_not_http_cached_and_recovers(self): |
| payload = gzip.compress(json.dumps({"v": 1, "f": {"ready": {"s": 4}}}).encode()) |
| session = FakeSession([FakeResponse(503), FakeResponse(200, chunks=(payload,))]) |
| app.app.state.http_session = session |
| app.app.state.reader_assets_lock = asyncio.Lock() |
| with patch.object(app, "reader_assets_cache", {"v": 1, "f": {}}), \ |
| patch.object(app, "reader_assets_cache_at", 0), \ |
| patch.object(app.time, "monotonic", return_value=10000) as clock: |
| failed = await app.api_reader_assets() |
| self.assertEqual(failed.headers["cache-control"], "no-store") |
| await app.get_reader_assets() |
| self.assertEqual(len(session.urls), 1) |
| clock.return_value += 6 |
| recovered = await app.api_reader_assets() |
| self.assertEqual(json.loads(recovered.body)["f"], {"ready": {"s": 4}}) |
| self.assertEqual(len(session.urls), 2) |
|
|
| async def test_reader_bucket_cache_hit_preserves_success_cache_control(self): |
| path = f"objects/aa/{'a' * 64}/{'b' * 16}/page-manifest.json" |
| with patch.object(app, "read_reader_bucket_resource", return_value=b"manifest") as read: |
| first = await app.api_reader_bucket_resource(path) |
| second = await app.api_reader_bucket_resource(path) |
| self.assertEqual(first.body, b"manifest") |
| self.assertEqual(second.body, b"manifest") |
| self.assertEqual(read.call_count, 1) |
| self.assertEqual(second.headers["cache-control"], "public, max-age=3600") |
|
|
| async def test_reader_bucket_page_uses_immutable_cache_control(self): |
| path = f"objects/aa/{'a' * 64}/{'b' * 16}/pages/page-000002.webp" |
| with patch.object(app, "read_reader_bucket_resource", return_value=b"page"): |
| response = await app.api_reader_bucket_resource(path) |
| self.assertEqual(response.headers["cache-control"], "public, max-age=31536000, immutable") |
|
|
| async def test_reader_bucket_single_flight_coalesces_uncached_page(self): |
| path = f"objects/aa/{'a' * 64}/{'b' * 16}/pages/page-000002.webp" |
| started = threading.Event() |
| release = threading.Event() |
|
|
| def read(_path): |
| started.set() |
| release.wait(1) |
| return b"page" |
|
|
| with patch.object(app, "read_reader_bucket_resource", side_effect=read) as bucket_read: |
| first = asyncio.create_task(app.api_reader_bucket_resource(path)) |
| self.assertTrue(await asyncio.to_thread(started.wait, 1)) |
| second = asyncio.create_task(app.api_reader_bucket_resource(path)) |
| await asyncio.sleep(0) |
| self.assertEqual(bucket_read.call_count, 1) |
| release.set() |
| responses = await asyncio.gather(first, second) |
| self.assertEqual([response.body for response in responses], [b"page", b"page"]) |
| self.assertEqual(bucket_read.call_count, 1) |
| self.assertNotIn(path, app.app.state.reader_bucket_cache) |
| self.assertEqual(app.app.state.reader_bucket_inflight, {}) |
|
|
| async def test_reader_bucket_semaphore_bounds_distinct_reads(self): |
| app.app.state.reader_bucket_semaphore = asyncio.Semaphore(1) |
| first_path = f"objects/aa/{'a' * 64}/{'b' * 16}/pages/page-000002.webp" |
| second_path = f"objects/aa/{'a' * 64}/{'b' * 16}/pages/page-000003.webp" |
| first_started = threading.Event() |
| release = threading.Event() |
|
|
| def read(path): |
| if path == first_path: |
| first_started.set() |
| release.wait(1) |
| return path.encode() |
|
|
| with patch.object(app, "read_reader_bucket_resource", side_effect=read) as bucket_read: |
| first = asyncio.create_task(app.api_reader_bucket_resource(first_path)) |
| self.assertTrue(await asyncio.to_thread(first_started.wait, 1)) |
| second = asyncio.create_task(app.api_reader_bucket_resource(second_path)) |
| await asyncio.sleep(0.05) |
| self.assertEqual(bucket_read.call_count, 1) |
| release.set() |
| await asyncio.gather(first, second) |
| self.assertEqual(bucket_read.call_count, 2) |
|
|
| async def test_reader_bucket_failures_are_no_store_and_not_cached(self): |
| path = f"objects/aa/{'a' * 64}/{'b' * 16}/page-manifest.json" |
| invalid = await app.api_reader_bucket_resource("invalid") |
| self.assertEqual(invalid.status_code, 400) |
| self.assertEqual(invalid.headers["cache-control"], "no-store") |
|
|
| for error, status in ((FileNotFoundError(), 404), (RuntimeError(), 502)): |
| with self.subTest(status=status), patch.object( |
| app, "read_reader_bucket_resource", side_effect=error |
| ) as bucket_read: |
| first = await app.api_reader_bucket_resource(path) |
| second = await app.api_reader_bucket_resource(path) |
| self.assertEqual(first.status_code, status) |
| self.assertEqual(second.status_code, status) |
| self.assertEqual(first.headers["cache-control"], "no-store") |
| self.assertEqual(bucket_read.call_count, 2) |
|
|
| async def test_reader_bucket_cache_evicts_by_bytes_and_skips_oversize_items(self): |
| manifest = f"objects/aa/{'a' * 64}/{'b' * 16}/page-manifest.json" |
| first_page = f"objects/aa/{'a' * 64}/{'b' * 16}/pages/page-000001.webp" |
| other_manifest = f"objects/bb/{'c' * 64}/{'d' * 16}/page-manifest.json" |
| old_total = app.READER_BUCKET_CACHE_MAX_BYTES |
| old_item = app.READER_BUCKET_CACHE_MAX_ITEM_BYTES |
| app.READER_BUCKET_CACHE_MAX_BYTES = 5 |
| app.READER_BUCKET_CACHE_MAX_ITEM_BYTES = 4 |
| payloads = {manifest: b"aaa", first_page: b"bbb", other_manifest: b"12345"} |
| try: |
| with patch.object(app, "read_reader_bucket_resource", side_effect=payloads.get) as bucket_read: |
| await app.api_reader_bucket_resource(manifest) |
| await app.api_reader_bucket_resource(first_page) |
| self.assertNotIn(manifest, app.app.state.reader_bucket_cache) |
| self.assertEqual(app.app.state.reader_bucket_cache_bytes, 3) |
| await app.api_reader_bucket_resource(manifest) |
| await app.api_reader_bucket_resource(other_manifest) |
| await app.api_reader_bucket_resource(other_manifest) |
| self.assertEqual(bucket_read.call_count, 5) |
| self.assertNotIn(other_manifest, app.app.state.reader_bucket_cache) |
| self.assertLessEqual(app.app.state.reader_bucket_cache_bytes, 5) |
| finally: |
| app.READER_BUCKET_CACHE_MAX_BYTES = old_total |
| app.READER_BUCKET_CACHE_MAX_ITEM_BYTES = old_item |
|
|
| async def test_reader_bucket_shutdown_cancels_flights_and_clears_cache(self): |
| task = asyncio.create_task(asyncio.sleep(10)) |
| app.app.state.reader_bucket_inflight["path"] = task |
| app.app.state.reader_bucket_cache["path"] = (time.monotonic(), b"data") |
| app.app.state.reader_bucket_cache_bytes = 4 |
| await app.shutdown_reader_bucket_state(app.app) |
| self.assertTrue(task.cancelled()) |
| self.assertEqual(app.app.state.reader_bucket_inflight, {}) |
| self.assertEqual(app.app.state.reader_bucket_cache, {}) |
| self.assertEqual(app.app.state.reader_bucket_cache_bytes, 0) |
|
|
| async def test_relative_redirect_is_resolved_against_current_url(self): |
| session = FakeSession([FakeResponse(302, {"Location": "next"}), FakeResponse(200)]) |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf" |
| result = await app.open_download_response( |
| session, start, aiohttp.ClientTimeout(total=1), |
| url_validator=app.validate_voiceofml_source_url, |
| ) |
| self.assertEqual(result.status, 200) |
| self.assertEqual([url for url, _ in session.urls], [start, "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/next"]) |
|
|
| async def test_redirect_cannot_escape_voiceofml_dataset_scope(self): |
| first = FakeResponse(302, {"Location": "https://huggingface.co/datasets/Other/Test/resolve/main/book.pdf"}) |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf" |
| with self.assertRaises(ValueError): |
| await app.open_download_response( |
| FakeSession([first]), start, aiohttp.ClientTimeout(total=1), |
| url_validator=app.validate_voiceofml_source_url, |
| ) |
| self.assertTrue(first.released) |
|
|
| async def test_redirect_cannot_escape_reader_assets_scope(self): |
| digest = "b" * 64 |
| start = f"https://huggingface.co/datasets/vomebook/Reader-Assets/resolve/main/objects/aa/{digest}/document.pdf" |
| first = FakeResponse(302, {"Location": "/datasets/VoiceOfML/Test/resolve/main/book.pdf"}) |
| with self.assertRaises(ValueError): |
| await app.open_download_response( |
| FakeSession([first]), start, aiohttp.ClientTimeout(total=1), |
| url_validator=app.validate_reader_source_url, |
| ) |
| self.assertTrue(first.released) |
|
|
| async def test_multilayer_encoded_reader_assets_redirect_cannot_escape_scope(self): |
| digest = "b" * 64 |
| start = f"https://huggingface.co/datasets/%252576omebook/Reader-Assets/resolve/main/objects/aa/{digest}/document.pdf" |
| first = FakeResponse(302, {"Location": "/datasets/VoiceOfML/Test/resolve/main/book.pdf"}) |
| self.assertEqual(app.validate_reader_source_url(start), start) |
| with self.assertRaises(ValueError): |
| await app.open_download_response( |
| FakeSession([first, FakeResponse(200)]), start, aiohttp.ClientTimeout(total=1), |
| url_validator=app.validate_reader_source_url, |
| ) |
| self.assertTrue(first.released) |
|
|
| async def test_redirect_to_disallowed_host_is_rejected_and_releases_first(self): |
| first = FakeResponse(302, {"Location": "https://evil.test/x"}) |
| with self.assertRaises(ValueError): |
| await app.open_download_response(FakeSession([first]), "https://huggingface.co/x", aiohttp.ClientTimeout(total=1)) |
| self.assertTrue(first.released) |
|
|
| async def test_redirect_without_location_is_rejected_and_released(self): |
| response = FakeResponse(302) |
| with self.assertRaises(ValueError): |
| await app.open_download_response(FakeSession([response]), "https://huggingface.co/x", aiohttp.ClientTimeout(total=1)) |
| self.assertTrue(response.released) |
|
|
| async def test_too_many_redirects_are_rejected(self): |
| responses = [FakeResponse(302, {"Location": "/next"}) for _ in range(app.MAX_REDIRECTS + 1)] |
| with self.assertRaises(ValueError): |
| await app.open_download_response(FakeSession(responses), "https://huggingface.co/x", aiohttp.ClientTimeout(total=1)) |
|
|
| async def test_resolved_redirect_target_is_reused_and_scoped(self): |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/cached.pdf" |
| final = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/actual.pdf" |
| first = FakeSession([FakeResponse(302, {"Location": "actual.pdf"}), FakeResponse(200)]) |
| result = await app.open_download_response( |
| first, start, aiohttp.ClientTimeout(total=1), |
| url_validator=app.validate_voiceofml_source_url, |
| ) |
| self.assertEqual(result.status, 200) |
| self.assertEqual([url for url, _ in first.urls], [start, final]) |
| second = FakeSession([FakeResponse(200)]) |
| result = await app.open_download_response( |
| second, start, aiohttp.ClientTimeout(total=1), |
| url_validator=app.validate_voiceofml_source_url, |
| ) |
| self.assertEqual(result.status, 200) |
| self.assertEqual([url for url, _ in second.urls], [final]) |
|
|
| async def test_rejected_redirect_target_falls_back_to_fresh_walk(self): |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/retry.pdf" |
| stale = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/stale.pdf" |
| fresh = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/fresh.pdf" |
| upstream_http.upstream_redirect_cache[("GET", start)] = (time.monotonic(), stale) |
| session = FakeSession([ |
| FakeResponse(403), FakeResponse(302, {"Location": "fresh.pdf"}), FakeResponse(200), |
| ]) |
| result = await app.open_download_response( |
| session, start, aiohttp.ClientTimeout(total=1), |
| url_validator=app.validate_voiceofml_source_url, |
| ) |
| self.assertEqual(result.status, 200) |
| self.assertEqual([url for url, _ in session.urls], [stale, start, fresh]) |
| self.assertEqual(upstream_http.upstream_redirect_cache[("GET", start)][1], fresh) |
|
|
| async def test_out_of_scope_cached_target_is_ignored(self): |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/scope.pdf" |
| escaped = "https://huggingface.co/datasets/vomebook/Reader-Assets/resolve/main/objects/aa/" + "b" * 64 + "/document.pdf" |
| upstream_http.upstream_redirect_cache[("GET", start)] = (time.monotonic(), escaped) |
| session = FakeSession([FakeResponse(200)]) |
| result = await app.open_download_response( |
| session, start, aiohttp.ClientTimeout(total=1), |
| ) |
| self.assertEqual(result.status, 200) |
| self.assertEqual([url for url, _ in session.urls], [start]) |
|
|
| async def test_non_redirect_response_is_returned_without_release(self): |
| response = FakeResponse(404) |
| result = await app.open_download_response(FakeSession([response]), "https://huggingface.co/x", aiohttp.ClientTimeout(total=1)) |
| self.assertIs(result, response) |
| self.assertFalse(response.released) |
|
|
| async def test_download_rejects_external_source_before_session_call(self): |
| session = FakeSession([]) |
| app.app.state.http_session = session |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) |
| result = await app.api_download("x.pdf", "https://evil.test/x") |
| self.assertEqual(result.status_code, 403) |
| self.assertEqual(session.urls, []) |
|
|
| async def test_download_upstream_error_releases_semaphore(self): |
| session = FakeSession([FakeResponse(502)]) |
| semaphore = asyncio.Semaphore(1) |
| app.app.state.http_session = session |
| app.app.state.upstream_semaphore = semaphore |
| result = await app.api_download("x.pdf", "https://huggingface.co/datasets/VoiceOfML/A/x") |
| self.assertEqual(result.status_code, 502) |
| self.assertEqual(semaphore._value, 1) |
|
|
| async def test_download_cancellation_releases_semaphore(self): |
| semaphore = asyncio.Semaphore(1) |
| app.app.state.http_session = FakeSession([]) |
| app.app.state.upstream_semaphore = semaphore |
| with patch.object(app, "open_download_response", AsyncMock(side_effect=asyncio.CancelledError)): |
| with self.assertRaises(asyncio.CancelledError): |
| await app.api_download("x.pdf", "https://huggingface.co/datasets/VoiceOfML/A/x") |
| self.assertEqual(semaphore._value, 1) |
|
|
| async def test_download_non_200_releases_upstream_and_semaphore(self): |
| upstream = FakeResponse(404) |
| app.app.state.http_session = FakeSession([upstream]) |
| semaphore = asyncio.Semaphore(1) |
| app.app.state.upstream_semaphore = semaphore |
| result = await app.api_download("x.pdf", "https://huggingface.co/datasets/VoiceOfML/A/x") |
| self.assertEqual(result.status_code, 502) |
| self.assertTrue(upstream.released) |
| self.assertEqual(semaphore._value, 1) |
|
|
| async def test_download_response_stream_releases_resources_after_consumption(self): |
| upstream = FakeResponse(200, chunks=(b"a", b"b")) |
| app.app.state.http_session = FakeSession([upstream]) |
| semaphore = asyncio.Semaphore(1) |
| app.app.state.upstream_semaphore = semaphore |
| response = await app.api_download("../book.pdf", "https://huggingface.co/datasets/VoiceOfML/A/x") |
| chunks = [chunk async for chunk in response.body_iterator] |
| self.assertEqual(chunks, [b"a", b"b"]) |
| self.assertTrue(upstream.released) |
| self.assertEqual(semaphore._value, 1) |
| self.assertIn("filename*=UTF-8''book.pdf", response.headers["content-disposition"]) |
|
|
| async def test_reader_content_forwards_range_and_preserves_partial_headers(self): |
| upstream = FakeResponse(206, { |
| "Content-Type": "application/pdf", |
| "Content-Length": "2", |
| "Content-Range": "bytes 0-1/10", |
| "Accept-Ranges": "bytes", |
| "ETag": '"reader"', |
| "Set-Cookie": "secret=1", |
| }, chunks=(b"ab",)) |
| session = FakeSession([upstream]) |
| semaphore = asyncio.Semaphore(1) |
| app.app.state.http_session = session |
| app.app.state.upstream_semaphore = semaphore |
| request = SimpleNamespace(method="GET", headers={ |
| "Range": "bytes=0-1", "If-None-Match": '"reader"', |
| "If-Modified-Since": "Tue, 25 Aug 2026 00:00:00 GMT", "Cookie": "secret=1", |
| }) |
| url = "https://huggingface.co/datasets/vomebook/Reader-Assets/resolve/main/objects/aa/" + "b" * 64 + "/document.pdf" |
| response = await app.api_reader_content(request, url) |
| self.assertEqual(response.status_code, 206) |
| self.assertEqual([chunk async for chunk in response.body_iterator], [b"ab"]) |
| self.assertEqual(session.urls[0][1]["headers"], { |
| "Accept-Encoding": "identity", "Range": "bytes=0-1", |
| "If-None-Match": '"reader"', "If-Modified-Since": "Tue, 25 Aug 2026 00:00:00 GMT", |
| }) |
| self.assertEqual(response.headers["content-range"], "bytes 0-1/10") |
| self.assertEqual(response.headers["content-length"], "2") |
| self.assertEqual(response.headers["content-encoding"], "identity") |
| self.assertIs(session.urls[0][1]["auto_decompress"], False) |
| self.assertNotIn("set-cookie", response.headers) |
| self.assertTrue(upstream.released) |
| self.assertEqual(semaphore._value, 1) |
|
|
| async def test_reader_content_reuses_completed_small_range_from_cache(self): |
| upstream = FakeResponse(206, { |
| "Content-Type": "application/pdf", "Content-Length": "2", |
| "Content-Range": "bytes 0-1/10", "Accept-Ranges": "bytes", |
| }, chunks=(b"ab",)) |
| session = FakeSession([upstream]) |
| app.app.state.http_session = session |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) |
| app.app.state.reader_range_cache = OrderedDict() |
| app.app.state.reader_range_cache_bytes = 0 |
| request = SimpleNamespace(method="GET", headers={"Range": "bytes=0-1"}) |
| url = "https://huggingface.co/datasets/vomebook/Reader-Assets/resolve/main/objects/aa/" + "b" * 64 + "/document.pdf" |
| first = await app.api_reader_content(request, url) |
| self.assertEqual([chunk async for chunk in first.body_iterator], [b"ab"]) |
| second = await app.api_reader_content(request, url) |
| self.assertEqual(second.body, b"ab") |
| self.assertEqual(len(session.urls), 1) |
|
|
| async def test_reader_content_does_not_reuse_range_cache_for_conditional_request(self): |
| first_upstream = FakeResponse(206, { |
| "Content-Type": "application/pdf", "Content-Length": "2", |
| "Content-Range": "bytes 0-1/10", "Accept-Ranges": "bytes", |
| "ETag": '"reader"', |
| }, chunks=(b"ab",)) |
| second_upstream = FakeResponse(304, {"ETag": '"reader"'}) |
| session = FakeSession([first_upstream, second_upstream]) |
| app.app.state.http_session = session |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) |
| app.app.state.reader_range_cache = OrderedDict() |
| app.app.state.reader_range_cache_bytes = 0 |
| url = "https://huggingface.co/datasets/vomebook/Reader-Assets/resolve/main/objects/aa/" + "b" * 64 + "/document.pdf" |
| first = await app.api_reader_content(SimpleNamespace(method="GET", headers={"Range": "bytes=0-1"}), url) |
| self.assertEqual([chunk async for chunk in first.body_iterator], [b"ab"]) |
| second = await app.api_reader_content(SimpleNamespace(method="GET", headers={ |
| "Range": "bytes=0-1", "If-None-Match": '"reader"', |
| }), url) |
| self.assertEqual(second.status_code, 304) |
| self.assertEqual(len(session.urls), 2) |
|
|
| async def test_reader_range_cache_does_not_turn_origin_200_into_206(self): |
| upstream = FakeResponse(200, {"Content-Type": "application/pdf", "Content-Length": "2"}, chunks=(b"ab",)) |
| session = FakeSession([upstream]) |
| app.app.state.http_session = session |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) |
| app.app.state.reader_range_cache = OrderedDict() |
| app.app.state.reader_range_cache_bytes = 0 |
| url = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf" |
| response = await app.api_reader_content(SimpleNamespace(method="GET", headers={"Range": "bytes=0-1"}), url) |
| self.assertEqual(response.status_code, 200) |
| self.assertEqual([chunk async for chunk in response.body_iterator], [b"ab"]) |
| self.assertEqual(app.app.state.reader_range_cache, OrderedDict()) |
|
|
| async def test_reader_content_rejects_oversized_content_length_before_streaming(self): |
| upstream = FakeResponse(200, {"Content-Length": str(app.READER_DOCUMENT_MAX_BYTES + 1)}, chunks=(b"not-read",)) |
| semaphore = asyncio.Semaphore(1) |
| app.app.state.http_session = FakeSession([upstream]) |
| app.app.state.upstream_semaphore = semaphore |
| response = await app.api_reader_content( |
| SimpleNamespace(method="GET", headers={}), |
| "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf", |
| ) |
| self.assertEqual(response.status_code, 413) |
| self.assertTrue(upstream.released) |
| self.assertEqual(semaphore._value, 1) |
|
|
| async def test_reader_content_rejects_oversized_range_total(self): |
| old_limit = app.READER_DOCUMENT_MAX_BYTES |
| app.READER_DOCUMENT_MAX_BYTES = 3 |
| upstream = FakeResponse(206, { |
| "Content-Length": "1", |
| "Content-Range": "bytes 0-0/4", |
| }, chunks=(b"a",)) |
| semaphore = asyncio.Semaphore(1) |
| app.app.state.http_session = FakeSession([upstream]) |
| app.app.state.upstream_semaphore = semaphore |
| try: |
| response = await app.api_reader_content( |
| SimpleNamespace(method="GET", headers={}), |
| "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf", |
| ) |
| finally: |
| app.READER_DOCUMENT_MAX_BYTES = old_limit |
| self.assertEqual(response.status_code, 413) |
| self.assertTrue(upstream.released) |
| self.assertEqual(semaphore._value, 1) |
|
|
| async def test_reader_content_enforces_actual_stream_limit_and_releases(self): |
| old_limit = app.READER_DOCUMENT_MAX_BYTES |
| app.READER_DOCUMENT_MAX_BYTES = 3 |
| upstream = FakeResponse(200, {}, chunks=(b"ab", b"cd")) |
| semaphore = asyncio.Semaphore(1) |
| app.app.state.http_session = FakeSession([upstream]) |
| app.app.state.upstream_semaphore = semaphore |
| try: |
| response = await app.api_reader_content( |
| SimpleNamespace(method="GET", headers={}), |
| "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf", |
| ) |
| with self.assertRaises(app.ReaderDocumentLimitError): |
| [chunk async for chunk in response.body_iterator] |
| finally: |
| app.READER_DOCUMENT_MAX_BYTES = old_limit |
| self.assertTrue(upstream.released) |
| self.assertEqual(semaphore._value, 1) |
|
|
| async def test_reader_content_cancellation_releases_semaphore(self): |
| semaphore = asyncio.Semaphore(1) |
| app.app.state.upstream_semaphore = semaphore |
| request = SimpleNamespace(method="GET", headers={}) |
| with patch.object(app, "open_download_response", AsyncMock(side_effect=asyncio.CancelledError)): |
| with self.assertRaises(asyncio.CancelledError): |
| await app.api_reader_content( |
| request, "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf" |
| ) |
| self.assertEqual(semaphore._value, 1) |
|
|
| async def test_reader_content_head_releases_resources_without_streaming(self): |
| upstream = FakeResponse(200, {"Content-Length": "10", "Accept-Ranges": "bytes"}) |
| session = FakeSession([upstream]) |
| semaphore = asyncio.Semaphore(1) |
| app.app.state.http_session = session |
| app.app.state.upstream_semaphore = semaphore |
| request = SimpleNamespace(method="HEAD", headers={}) |
| response = await app.api_reader_content(request, "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.epub") |
| self.assertEqual(response.status_code, 200) |
| self.assertEqual(session.urls[0][1]["method"], "HEAD") |
| self.assertIs(session.urls[0][1]["auto_decompress"], False) |
| self.assertEqual(response.headers["content-length"], "10") |
| self.assertTrue(upstream.released) |
| self.assertEqual(semaphore._value, 1) |
|
|
| async def test_reader_content_rejects_external_source_before_session_call(self): |
| session = FakeSession([]) |
| app.app.state.http_session = session |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) |
| result = await app.api_reader_content(SimpleNamespace(method="GET", headers={}), "https://evil.test/book.pdf") |
| self.assertEqual(result.status_code, 403) |
| self.assertEqual(session.urls, []) |
|
|
| async def test_download_check_upstream_error_returns_502_and_releases(self): |
| upstream = FakeResponse(503) |
| app.app.state.http_session = FakeSession([upstream]) |
| result = await app.api_download_check("https://huggingface.co/datasets/VoiceOfML/A/x") |
| self.assertEqual(result.status_code, 502) |
| self.assertTrue(upstream.released) |
|
|
| async def test_download_check_uses_head_and_shared_semaphore(self): |
| upstream = FakeResponse(200) |
| session = FakeSession([upstream]) |
| semaphore = asyncio.Semaphore(1) |
| app.app.state.http_session = session |
| app.app.state.upstream_semaphore = semaphore |
| result = await app.api_download_check("https://huggingface.co/datasets/VoiceOfML/A/x") |
| self.assertEqual(result.status_code, 200) |
| self.assertEqual(session.urls[0][1]["method"], "HEAD") |
| self.assertTrue(upstream.released) |
| self.assertEqual(semaphore._value, 1) |
|
|
| async def test_local_txt_file_is_served_without_remote_request(self): |
| with tempfile.TemporaryDirectory() as directory: |
| old_dir = app.TXT_DIR |
| app.TXT_DIR = Path(directory) |
| local = Path(directory) / "local.txt" |
| local.write_text("local", encoding="utf-8") |
| session = FakeSession([]) |
| app.app.state.http_session = session |
| try: |
| response = await app.serve_txt_proxy("local.txt") |
| self.assertEqual(response.path, local) |
| self.assertEqual(session.urls, []) |
| finally: |
| app.TXT_DIR = old_dir |
|
|
| async def test_remote_txt_success_streams_content(self): |
| upstream = FakeResponse(200, chunks=(b"remote",)) |
| app.app.state.http_session = FakeSession([upstream]) |
| with tempfile.TemporaryDirectory() as directory, patch.object(app, "resolve_safe_txt_path", return_value=Path(directory) / "missing.txt"): |
| response = await app.serve_txt_proxy("remote.txt") |
| self.assertEqual([chunk async for chunk in response.body_iterator], [b"remote"]) |
| self.assertTrue(upstream.released) |
|
|
| async def test_txt_traversal_is_400_without_remote_request(self): |
| session = FakeSession([]) |
| app.app.state.http_session = session |
| response = await app.serve_txt_proxy("../secret.txt") |
| self.assertEqual(response.status_code, 400) |
| self.assertEqual(session.urls, []) |
|
|
| async def test_remote_txt_404_maps_to_404(self): |
| upstream = FakeResponse(404) |
| app.app.state.http_session = FakeSession([upstream]) |
| with tempfile.TemporaryDirectory() as directory, patch.object(app, "resolve_safe_txt_path", return_value=Path(directory) / "missing.txt"): |
| response = await app.serve_txt_proxy("missing.txt") |
| self.assertEqual(response.status_code, 404) |
| self.assertTrue(upstream.released) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
