mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 10:15:21 +00:00
fix(api): address review on the event-loop streams
- Per-user SSE cap uses per-connection leases in a sorted set instead of a shared INCR/DECR counter with a TTL. A stream that outlived the TTL could let the counter expire, then decrement another stream's slot or drive the count negative. Leases refresh while a stream sends frames and age out when a stream dies without cleanup. - Shielded cleanup is bounded per step: stream close and on_close in ClosingStreamingResponse, unsubscribe and close in AsyncTopic, lease release, and the replay-budget check, so a dead Redis connection can't hold a request or a graceful shutdown. - Artifact downloads send an ASCII filename plus an RFC 5987 filename* for non-ASCII names, build headers before opening the file, and close disk handles off the event loop. - Type hints and docstrings on the new helpers.
This commit is contained in:
1 parent
54b540ea42
commit
e889392eba
17 files changed
+911
-255
No files matched your search
@@ -0,0 +1,92 @@
|
||||
"""Tests for ``docsgpt/api/asgi_stream.py`` — cleanup that can't stall a worker.
|
||||
|
||||
The cleanup runs shielded from the cancellation that ended the response, so
|
||||
each step needs its own deadline: a close or release waiting on a dead Redis
|
||||
connection must not hold the request task, or a graceful shutdown, forever.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from docsgpt.api import asgi_stream
|
||||
from docsgpt.api.asgi_stream import ClosingStreamingResponse
|
||||
from tests.asgi_stream import stream_then_disconnect
|
||||
|
||||
|
||||
class _StuckCloseIterator:
|
||||
"""Yields one chunk, then ends; its ``aclose`` never returns."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._sent = False
|
||||
self.aclose_started = False
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._sent:
|
||||
raise StopAsyncIteration
|
||||
self._sent = True
|
||||
return b"chunk"
|
||||
|
||||
async def aclose(self):
|
||||
self.aclose_started = True
|
||||
await anyio.sleep_forever()
|
||||
|
||||
|
||||
async def _one_chunk():
|
||||
yield b"chunk"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.asyncio
|
||||
class TestBoundedCleanup:
|
||||
async def test_stuck_stream_close_is_abandoned_and_on_close_still_runs(self, monkeypatch):
|
||||
monkeypatch.setattr(asgi_stream, "_CLEANUP_TIMEOUT_SECONDS", 0.05)
|
||||
body = _StuckCloseIterator()
|
||||
released = {"ran": False}
|
||||
|
||||
async def on_close():
|
||||
released["ran"] = True
|
||||
|
||||
response = ClosingStreamingResponse(body, media_type="text/plain", on_close=on_close)
|
||||
with anyio.fail_after(2):
|
||||
status, _, chunks = await stream_then_disconnect(
|
||||
response, "/", chunks_before_disconnect=10**6
|
||||
)
|
||||
assert status == 200
|
||||
assert chunks == [b"chunk"]
|
||||
assert body.aclose_started
|
||||
assert released["ran"] is True
|
||||
|
||||
async def test_stuck_on_close_is_abandoned(self, monkeypatch):
|
||||
monkeypatch.setattr(asgi_stream, "_CLEANUP_TIMEOUT_SECONDS", 0.05)
|
||||
|
||||
async def on_close():
|
||||
await anyio.sleep_forever()
|
||||
|
||||
response = ClosingStreamingResponse(_one_chunk(), media_type="text/plain", on_close=on_close)
|
||||
with anyio.fail_after(2):
|
||||
status, _, chunks = await stream_then_disconnect(
|
||||
response, "/", chunks_before_disconnect=10**6
|
||||
)
|
||||
assert status == 200
|
||||
assert chunks == [b"chunk"]
|
||||
|
||||
async def test_on_close_runs_when_client_leaves_before_first_frame(self):
|
||||
released = {"count": 0}
|
||||
|
||||
async def on_close():
|
||||
released["count"] += 1
|
||||
|
||||
async def endless():
|
||||
while True:
|
||||
await anyio.sleep(0.01)
|
||||
yield b"tick"
|
||||
|
||||
response = ClosingStreamingResponse(endless(), media_type="text/plain", on_close=on_close)
|
||||
with anyio.fail_after(2):
|
||||
await stream_then_disconnect(response, "/", chunks_before_disconnect=0)
|
||||
assert released["count"] == 1
|
||||
@@ -1,26 +1,29 @@
|
||||
"""Tests for ``docsgpt/api/async_sse.py``.
|
||||
|
||||
Native-async reconnect endpoint: GET /api/messages/<id>/events. Auth gate,
|
||||
ownership gate, malformed-id rejection, Last-Event-ID normalisation, and the
|
||||
SSE response shape (headers + ``: connected`` prelude). The route is a
|
||||
Starlette endpoint, so it's driven through Starlette's TestClient over a
|
||||
minimal app built from ``async_sse_routes``.
|
||||
ownership gate, malformed-id rejection, Last-Event-ID normalisation, the
|
||||
per-user connection lease, and the SSE response shape (headers + ``: connected``
|
||||
prelude). The route is a Starlette endpoint, so it's driven through Starlette's
|
||||
TestClient over a minimal app built from ``async_sse_routes``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from starlette.applications import Starlette
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from docsgpt.api import async_sse
|
||||
from docsgpt.api.async_sse import (
|
||||
_MESSAGE_ID_RE,
|
||||
_normalise_last_event_id,
|
||||
async_sse_routes,
|
||||
)
|
||||
from docsgpt.core.settings import settings
|
||||
from tests.asgi_stream import stream_then_disconnect
|
||||
from tests.fake_async_redis import FakeAsyncRedis
|
||||
|
||||
VALID_UUID = "67d65e8f-e7fb-4df1-9e6e-99ea6c830206"
|
||||
|
||||
@@ -28,42 +31,54 @@ _AUTH = "docsgpt.api.asgi_auth.handle_auth"
|
||||
_OWNS = "docsgpt.api.async_sse._user_owns_message"
|
||||
_STREAM = "docsgpt.api.async_sse.build_message_event_stream_async"
|
||||
_AREDIS = "docsgpt.api.async_sse.get_async_redis_instance"
|
||||
_LEASES = "user:alice:sse_leases"
|
||||
|
||||
|
||||
def _app() -> Starlette:
|
||||
return Starlette(routes=async_sse_routes)
|
||||
|
||||
|
||||
def _client() -> TestClient:
|
||||
return TestClient(Starlette(routes=async_sse_routes))
|
||||
return TestClient(_app())
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_redis_by_default():
|
||||
def _no_redis_by_default(monkeypatch):
|
||||
"""Disable the per-user cap (no Redis) so gate tests stay hermetic.
|
||||
|
||||
Cap-specific tests override this with their own mock Redis.
|
||||
Cap-specific tests override this with their own fake Redis.
|
||||
"""
|
||||
monkeypatch.setattr(async_sse.settings, "SSE_MAX_CONCURRENT_PER_USER", 8)
|
||||
monkeypatch.setattr(async_sse.settings, "SSE_KEEPALIVE_SECONDS", 15)
|
||||
with patch(_AREDIS, AsyncMock(return_value=None)):
|
||||
yield
|
||||
|
||||
|
||||
def _mock_redis(incr_value: int) -> AsyncMock:
|
||||
redis = AsyncMock()
|
||||
redis.incr = AsyncMock(return_value=incr_value)
|
||||
redis.expire = AsyncMock(return_value=True)
|
||||
redis.decr = AsyncMock(return_value=incr_value - 1)
|
||||
return redis
|
||||
|
||||
|
||||
def _fake_stream(record: dict | None = None):
|
||||
"""Async builder stub: records the cursor, yields the prelude, returns."""
|
||||
def _fake_stream(record: dict | None = None, redis: FakeAsyncRedis | None = None):
|
||||
"""Async builder stub: records the cursor (and live leases), yields the prelude, returns."""
|
||||
|
||||
async def _gen(message_id, last_event_id=None, **kwargs):
|
||||
if record is not None:
|
||||
record["message_id"] = message_id
|
||||
record["last_event_id"] = last_event_id
|
||||
if redis is not None:
|
||||
record["leases_during"] = len(redis.zsets.get(_LEASES, {}))
|
||||
yield ": connected\n\n"
|
||||
|
||||
return _gen
|
||||
|
||||
|
||||
def _endless_stream():
|
||||
async def _gen(message_id, last_event_id=None, **kwargs):
|
||||
import anyio
|
||||
|
||||
while True:
|
||||
await anyio.sleep(0.01)
|
||||
yield ": keepalive\n\n"
|
||||
|
||||
return _gen
|
||||
|
||||
|
||||
# ── pure helpers ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@@ -206,17 +221,14 @@ class TestLastEventIdParsing:
|
||||
assert captured["last_event_id"] is None
|
||||
|
||||
|
||||
# ── per-user concurrent-connection cap ──────────────────────────────────────
|
||||
# ── per-user connection lease ───────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestConnectionCap:
|
||||
def _cap(self) -> int:
|
||||
return int(settings.SSE_MAX_CONCURRENT_PER_USER) or 8
|
||||
|
||||
def test_429_when_over_cap(self):
|
||||
cap = self._cap()
|
||||
redis = _mock_redis(incr_value=cap + 1) # post-incr count exceeds cap
|
||||
def test_429_when_user_already_holds_the_cap(self):
|
||||
redis = FakeAsyncRedis()
|
||||
redis.zsets[_LEASES] = {f"tab-{i}": time.time() for i in range(8)}
|
||||
with patch(_AUTH, return_value={"sub": "alice"}), patch(
|
||||
_OWNS, return_value=True
|
||||
), patch(_STREAM, _fake_stream()), patch(
|
||||
@@ -224,23 +236,22 @@ class TestConnectionCap:
|
||||
):
|
||||
r = _client().get(f"/api/messages/{VALID_UUID}/events")
|
||||
assert r.status_code == 429
|
||||
# The increment is rolled back so a rejected attempt doesn't wedge
|
||||
# the counter at the cap forever.
|
||||
redis.decr.assert_awaited_once()
|
||||
# The rejected attempt removed its own lease and nobody else's.
|
||||
assert set(redis.zsets[_LEASES]) == {f"tab-{i}" for i in range(8)}
|
||||
|
||||
def test_200_and_slot_released_when_under_cap(self):
|
||||
redis = _mock_redis(incr_value=1)
|
||||
def test_200_lease_held_while_streaming_then_released(self):
|
||||
redis = FakeAsyncRedis()
|
||||
captured: dict = {}
|
||||
with patch(_AUTH, return_value={"sub": "alice"}), patch(
|
||||
_OWNS, return_value=True
|
||||
), patch(_STREAM, _fake_stream()), patch(
|
||||
), patch(_STREAM, _fake_stream(captured, redis)), patch(
|
||||
_AREDIS, AsyncMock(return_value=redis)
|
||||
):
|
||||
r = _client().get(f"/api/messages/{VALID_UUID}/events")
|
||||
assert r.status_code == 200
|
||||
assert ": connected" in r.text
|
||||
redis.incr.assert_awaited_once()
|
||||
# Slot released when the stream finishes (terminal/close).
|
||||
redis.decr.assert_awaited_once()
|
||||
assert captured["leases_during"] == 1
|
||||
assert redis.zsets == {}
|
||||
|
||||
def test_cap_skipped_when_redis_unavailable(self):
|
||||
# The autouse fixture already makes get_async_redis_instance -> None;
|
||||
@@ -250,3 +261,18 @@ class TestConnectionCap:
|
||||
), patch(_STREAM, _fake_stream()):
|
||||
r = _client().get(f"/api/messages/{VALID_UUID}/events")
|
||||
assert r.status_code == 200
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect_releases_lease(self):
|
||||
redis = FakeAsyncRedis()
|
||||
with patch(_AUTH, return_value={"sub": "alice"}), patch(
|
||||
_OWNS, return_value=True
|
||||
), patch(_STREAM, _endless_stream()), patch(
|
||||
_AREDIS, AsyncMock(return_value=redis)
|
||||
):
|
||||
status, _, chunks = await stream_then_disconnect(
|
||||
_app(), f"/api/messages/{VALID_UUID}/events", chunks_before_disconnect=1
|
||||
)
|
||||
assert status == 200
|
||||
assert chunks[0] == b": keepalive\n\n"
|
||||
assert redis.zsets == {}
|
||||
@@ -1,17 +1,20 @@
|
||||
"""Tests for docsgpt/api/events/routes.py — the native-async ``GET /api/events``.
|
||||
|
||||
The route is a Starlette endpoint served on the event loop. Tests drive it
|
||||
through Starlette's ``TestClient`` with the async Redis client and
|
||||
``AsyncTopic.subscribe`` faked. Fake subscriptions end after a few frames so
|
||||
the stream finishes and the client returns the body; a mid-stream client
|
||||
disconnect and the shutdown path go through ``tests/asgi_stream.py``.
|
||||
through Starlette's ``TestClient`` with the async Redis client faked
|
||||
(``tests/fake_async_redis.py``: real sorted sets for connection leases, mocks
|
||||
for streams) and ``AsyncTopic.subscribe`` replaced. Fake subscriptions end
|
||||
after a few frames so the stream finishes and the client returns the body; a
|
||||
mid-stream client disconnect and the shutdown path go through
|
||||
``tests/asgi_stream.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import json
|
||||
from typing import Any
|
||||
import time
|
||||
from typing import Any, Callable, Optional
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import anyio
|
||||
@@ -23,11 +26,12 @@ from docsgpt.api.events import routes as events_module
|
||||
from docsgpt.cache import PUBSUB_SOCKET_TIMEOUT_SECONDS
|
||||
from docsgpt.core.shutdown import begin_shutdown, reset_shutdown
|
||||
from tests.asgi_stream import stream_then_disconnect
|
||||
from tests.fake_async_redis import FakeAsyncRedis
|
||||
|
||||
_AUTH = "docsgpt.api.asgi_auth.handle_auth"
|
||||
_AREDIS = "docsgpt.api.events.routes.get_async_redis_instance"
|
||||
_SUBSCRIBE = "docsgpt.api.events.routes.AsyncTopic.subscribe"
|
||||
_COUNTER = "user:alice:sse_count"
|
||||
_LEASES = "user:alice:sse_leases"
|
||||
ALICE = {"sub": "alice"}
|
||||
OLD_CURSOR = "1735682300000-0"
|
||||
|
||||
@@ -40,22 +44,23 @@ def _get(headers: dict | None = None, params: dict | None = None):
|
||||
return TestClient(_app()).get("/api/events", headers=headers or {}, params=params or {})
|
||||
|
||||
|
||||
def _redis(incr: int = 1) -> AsyncMock:
|
||||
redis = AsyncMock()
|
||||
redis.incr = AsyncMock(return_value=incr)
|
||||
redis.expire = AsyncMock(return_value=True)
|
||||
redis.decr = AsyncMock(return_value=0)
|
||||
redis.xinfo_stream = AsyncMock(side_effect=Exception("no such key"))
|
||||
redis.xrange = AsyncMock(return_value=[])
|
||||
return redis
|
||||
def _redis() -> FakeAsyncRedis:
|
||||
return FakeAsyncRedis()
|
||||
|
||||
|
||||
def _subscribe(*payloads: Any, fire_callback: bool = True, record: dict | None = None):
|
||||
def _subscribe(
|
||||
*payloads: Any,
|
||||
fire_callback: bool = True,
|
||||
record: dict | None = None,
|
||||
probe: Optional[Callable[[], None]] = None,
|
||||
):
|
||||
"""``AsyncTopic.subscribe`` stand-in: ack, one idle tick, the payloads, then end."""
|
||||
|
||||
async def _impl(self, on_subscribe=None, poll_timeout=1.0, **kwargs):
|
||||
if record is not None:
|
||||
record.update(kwargs, poll_timeout=poll_timeout, topic=self.name)
|
||||
if probe is not None:
|
||||
probe()
|
||||
if fire_callback and on_subscribe is not None:
|
||||
result = on_subscribe()
|
||||
if inspect.isawaitable(result):
|
||||
@@ -180,6 +185,7 @@ class TestStreamShape:
|
||||
assert r.status_code == 200
|
||||
assert r.text == ": connected\n\n: push_disabled\n\n"
|
||||
redis.incr.assert_not_awaited()
|
||||
assert redis.zsets == {}
|
||||
|
||||
def test_subscribes_to_user_topic_with_liveness_probe(self):
|
||||
record: dict = {}
|
||||
@@ -200,22 +206,34 @@ class TestStreamShape:
|
||||
assert ": keepalive\n\n" in r.text
|
||||
|
||||
|
||||
# ── concurrency cap ─────────────────────────────────────────────────────
|
||||
# ── concurrency cap (per-connection leases) ─────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestConcurrencyCap:
|
||||
def test_429_when_over_cap_and_increment_rolled_back(self):
|
||||
redis = _redis(incr=9)
|
||||
def test_429_when_user_already_holds_the_cap(self):
|
||||
redis = _redis()
|
||||
redis.zsets[_LEASES] = {f"tab-{i}": time.time() for i in range(8)}
|
||||
with patch(_AUTH, return_value=ALICE), patch(
|
||||
_AREDIS, AsyncMock(return_value=redis)
|
||||
), patch(_SUBSCRIBE, _subscribe()):
|
||||
r = _get()
|
||||
assert r.status_code == 429
|
||||
assert r.json() == {"success": False, "message": "Too many concurrent SSE connections"}
|
||||
redis.decr.assert_awaited_once_with(_COUNTER)
|
||||
# The rejected attempt removed its own lease and nobody else's.
|
||||
assert set(redis.zsets[_LEASES]) == {f"tab-{i}" for i in range(8)}
|
||||
|
||||
def test_cap_disabled_leaves_counter_alone(self, monkeypatch):
|
||||
def test_leases_left_by_dead_streams_age_out(self):
|
||||
redis = _redis()
|
||||
redis.zsets[_LEASES] = {f"crashed-{i}": time.time() - 3600 for i in range(8)}
|
||||
with patch(_AUTH, return_value=ALICE), patch(
|
||||
_AREDIS, AsyncMock(return_value=redis)
|
||||
), patch(_SUBSCRIBE, _subscribe()):
|
||||
r = _get()
|
||||
assert r.status_code == 200
|
||||
assert redis.zsets == {}
|
||||
|
||||
def test_cap_disabled_takes_no_lease(self, monkeypatch):
|
||||
monkeypatch.setattr(events_module.settings, "SSE_MAX_CONCURRENT_PER_USER", 0)
|
||||
redis = _redis()
|
||||
with patch(_AUTH, return_value=ALICE), patch(
|
||||
@@ -223,38 +241,31 @@ class TestConcurrencyCap:
|
||||
), patch(_SUBSCRIBE, _subscribe()):
|
||||
r = _get()
|
||||
assert r.status_code == 200
|
||||
assert _COUNTER not in [call.args[0] for call in redis.incr.await_args_list]
|
||||
redis.decr.assert_not_awaited()
|
||||
assert redis.zsets == {}
|
||||
|
||||
def test_expire_failure_does_not_bypass_cap(self):
|
||||
redis = _redis(incr=9)
|
||||
redis.expire = AsyncMock(side_effect=Exception("expire blip"))
|
||||
with patch(_AUTH, return_value=ALICE), patch(
|
||||
_AREDIS, AsyncMock(return_value=redis)
|
||||
), patch(_SUBSCRIBE, _subscribe()):
|
||||
r = _get()
|
||||
assert r.status_code == 429
|
||||
|
||||
def test_incr_failure_fails_open(self):
|
||||
def test_redis_errors_fail_open(self):
|
||||
redis = _redis()
|
||||
redis.incr = AsyncMock(side_effect=Exception("redis down"))
|
||||
redis.fail = ConnectionError("reset by peer")
|
||||
with patch(_AUTH, return_value=ALICE), patch(
|
||||
_AREDIS, AsyncMock(return_value=redis)
|
||||
), patch(_SUBSCRIBE, _subscribe()):
|
||||
r = _get()
|
||||
assert r.status_code == 200
|
||||
# Nothing was counted, so nothing is released.
|
||||
redis.decr.assert_not_awaited()
|
||||
|
||||
def test_slot_released_when_stream_ends(self):
|
||||
redis = _redis(incr=1)
|
||||
def test_lease_held_while_streaming_and_released_at_end(self):
|
||||
redis = _redis()
|
||||
seen: dict = {}
|
||||
|
||||
def _count_leases():
|
||||
seen["during"] = len(redis.zsets.get(_LEASES, {}))
|
||||
|
||||
with patch(_AUTH, return_value=ALICE), patch(
|
||||
_AREDIS, AsyncMock(return_value=redis)
|
||||
), patch(_SUBSCRIBE, _subscribe()):
|
||||
), patch(_SUBSCRIBE, _subscribe(probe=_count_leases)):
|
||||
r = _get()
|
||||
assert r.status_code == 200
|
||||
redis.incr.assert_awaited_once_with(_COUNTER)
|
||||
redis.decr.assert_awaited_once_with(_COUNTER)
|
||||
assert seen["during"] == 1
|
||||
assert redis.zsets == {}
|
||||
|
||||
def test_redis_unavailable_serves_stream_without_cap(self):
|
||||
with patch(_AUTH, return_value=ALICE):
|
||||
@@ -269,9 +280,9 @@ class TestConcurrencyCap:
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.asyncio
|
||||
class TestStreamLifecycle:
|
||||
async def test_client_disconnect_releases_slot_and_closes_subscription(self, monkeypatch):
|
||||
async def test_client_disconnect_releases_lease_and_closes_subscription(self, monkeypatch):
|
||||
monkeypatch.setattr(events_module.settings, "SSE_KEEPALIVE_SECONDS", 0)
|
||||
redis = _redis(incr=1)
|
||||
redis = _redis()
|
||||
state: dict = {}
|
||||
with patch(_AUTH, return_value=ALICE), patch(
|
||||
_AREDIS, AsyncMock(return_value=redis)
|
||||
@@ -284,12 +295,12 @@ class TestStreamLifecycle:
|
||||
assert status == 200
|
||||
assert chunks[0] == b": connected\n\n"
|
||||
assert state.get("closed") is True
|
||||
redis.decr.assert_awaited_once_with(_COUNTER)
|
||||
assert redis.zsets == {}
|
||||
|
||||
async def test_disconnect_before_first_frame_still_releases_slot_once(self):
|
||||
async def test_disconnect_before_first_frame_still_releases_lease(self):
|
||||
# The tab closes as the response starts: the body generator may never
|
||||
# run, so its own cleanup can't be what frees the slot.
|
||||
redis = _redis(incr=1)
|
||||
redis = _redis()
|
||||
state: dict = {}
|
||||
with patch(_AUTH, return_value=ALICE), patch(
|
||||
_AREDIS, AsyncMock(return_value=redis)
|
||||
@@ -298,10 +309,10 @@ class TestStreamLifecycle:
|
||||
_app(), "/api/events", chunks_before_disconnect=0
|
||||
)
|
||||
assert status == 200
|
||||
redis.decr.assert_awaited_once_with(_COUNTER)
|
||||
assert redis.zsets == {}
|
||||
|
||||
async def test_shutdown_ends_stream_and_releases_slot(self):
|
||||
redis = _redis(incr=1)
|
||||
async def test_shutdown_ends_stream_and_releases_lease(self):
|
||||
redis = _redis()
|
||||
state: dict = {}
|
||||
begin_shutdown()
|
||||
with patch(_AUTH, return_value=ALICE), patch(
|
||||
@@ -313,7 +324,7 @@ class TestStreamLifecycle:
|
||||
assert status == 200
|
||||
assert b"".join(chunks) == b": connected\n\n"
|
||||
assert state.get("closed") is True
|
||||
redis.decr.assert_awaited_once_with(_COUNTER)
|
||||
assert redis.zsets == {}
|
||||
|
||||
|
||||
# ── replay + live tail ──────────────────────────────────────────────────
|
||||
@@ -457,11 +468,22 @@ class TestReplayBudget:
|
||||
assert redis.expire.await_count == 4
|
||||
assert all(call.args == ("user:alice:replay_count", 60) for call in redis.expire.await_args_list)
|
||||
|
||||
async def test_fail_open_on_redis_error(self, monkeypatch):
|
||||
async def test_fail_open_on_redis_error(self):
|
||||
redis = _redis()
|
||||
redis.incr = AsyncMock(side_effect=Exception("redis down"))
|
||||
assert await events_module._allow_replay(redis, "alice", OLD_CURSOR) is True
|
||||
|
||||
async def test_fail_open_when_redis_stalls(self, monkeypatch):
|
||||
monkeypatch.setattr(events_module, "_REDIS_OP_TIMEOUT_SECONDS", 0.05)
|
||||
redis = _redis()
|
||||
|
||||
async def _stalled(_key):
|
||||
await anyio.sleep_forever()
|
||||
|
||||
redis.incr = AsyncMock(side_effect=_stalled)
|
||||
with anyio.fail_after(2):
|
||||
assert await events_module._allow_replay(redis, "alice", OLD_CURSOR) is True
|
||||
|
||||
async def test_recovers_when_seeding_expire_raises(self):
|
||||
redis = _redis()
|
||||
counter = {"v": 0}
|
||||
@@ -522,19 +544,15 @@ class TestReplayBudgetRoute:
|
||||
live frames would advance the client's cursor past the un-replayed
|
||||
window. The 429 keeps the cursor pinned for the next attempt."""
|
||||
redis = _redis()
|
||||
|
||||
def _incr(key):
|
||||
return 1 if key == _COUNTER else 31
|
||||
|
||||
redis.incr = AsyncMock(side_effect=_incr)
|
||||
redis.incr = AsyncMock(return_value=31)
|
||||
with patch(_AUTH, return_value=ALICE), patch(
|
||||
_AREDIS, AsyncMock(return_value=redis)
|
||||
), patch(_SUBSCRIBE, _subscribe()):
|
||||
r = _get(headers={"Last-Event-ID": OLD_CURSOR})
|
||||
assert r.status_code == 429
|
||||
assert r.json() == {"success": False, "message": "Replay budget exhausted"}
|
||||
# The concurrency slot is released so a denied request doesn't hold it.
|
||||
redis.decr.assert_awaited_once_with(_COUNTER)
|
||||
# The connection lease is released so a denied request doesn't hold a slot.
|
||||
assert redis.zsets == {}
|
||||
|
||||
|
||||
# ── format helpers ──────────────────────────────────────────────────────
|
||||
|
||||
@@ -738,6 +738,58 @@ class TestDownloadArtifact:
|
||||
assert disposition == 'attachment; filename="aInjected: x.txt"'
|
||||
assert "injected" not in resp.headers
|
||||
|
||||
def test_unicode_filename_served_with_rfc5987_header(self, _patch_db, token_owner, monkeypatch):
|
||||
# Latin-1 header encoding used to reject this and return 400 after
|
||||
# opening the file.
|
||||
art = _seed_download(_patch_db, filename="报告.pdf", mime_type="application/pdf")
|
||||
storage = _use_storage(monkeypatch, _FakeStorage(b"PDF"))
|
||||
resp = _download(art["id"], token=token_owner)
|
||||
assert resp.status_code == 200
|
||||
assert resp.content == b"PDF"
|
||||
assert resp.headers["content-disposition"] == (
|
||||
f'attachment; filename="artifact-{art["id"]}.pdf"; '
|
||||
"filename*=UTF-8''%E6%8A%A5%E5%91%8A.pdf"
|
||||
)
|
||||
assert storage.opened[0].closed
|
||||
|
||||
def test_accented_filename_keeps_readable_ascii_fallback(self, _patch_db, token_owner, monkeypatch):
|
||||
art = _seed_download(_patch_db, filename="résumé final.pdf")
|
||||
_use_storage(monkeypatch, _FakeStorage(b"CV"))
|
||||
resp = _download(art["id"], token=token_owner)
|
||||
assert resp.status_code == 200
|
||||
assert resp.headers["content-disposition"] == (
|
||||
'attachment; filename="resume final.pdf"; '
|
||||
"filename*=UTF-8''r%C3%A9sum%C3%A9%20final.pdf"
|
||||
)
|
||||
|
||||
def test_disk_file_closed_off_the_event_loop(self, _patch_db, token_owner, monkeypatch):
|
||||
import threading
|
||||
|
||||
class _TrackedHandle:
|
||||
"""A non-BytesIO handle, like LocalStorage's open file."""
|
||||
|
||||
def __init__(self, data: bytes):
|
||||
self._buf = io.BytesIO(data)
|
||||
self.close_thread = None
|
||||
|
||||
def read(self, size=-1):
|
||||
return self._buf.read(size)
|
||||
|
||||
def close(self):
|
||||
self.close_thread = threading.current_thread().name
|
||||
|
||||
handle = _TrackedHandle(b"DISK")
|
||||
|
||||
class _HandleStorage(_FakeStorage):
|
||||
def get_file(self, path):
|
||||
return handle
|
||||
|
||||
art = _seed_download(_patch_db)
|
||||
_use_storage(monkeypatch, _HandleStorage())
|
||||
resp = _download(art["id"], token=token_owner)
|
||||
assert resp.content == b"DISK"
|
||||
assert handle.close_thread == "AnyIO worker thread"
|
||||
|
||||
def test_non_uuid_id_404(self, _patch_db, token_owner):
|
||||
resp = _download("legacy-mongo-objectid-aabbcc", token=token_owner)
|
||||
assert resp.status_code == 404
|
||||
|
||||
Reference in new issue
Block a user