chore: import upstream snapshot with attribution

python/assistant-stream/tests/test_resumable_in_memory.py

@@ -0,0 +1,308 @@
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from assistant_stream.resumable import (
+    ResumableStreamError,
+    create_in_memory_resumable_stream_store,
+)
+from assistant_stream.resumable.types import ResumableStreamEntry
+
+
+def _bytes(s: str) -> bytes:
+    return s.encode("utf-8")
+
+
+def _decode(chunk: bytes) -> str:
+    return chunk.decode("utf-8")
+
+
+async def _drain(iter_) -> list[str]:
+    out: list[str] = []
+    async for entry in iter_:
+        out.append(_decode(entry.chunk))
+    return out
+
+
+@pytest.mark.anyio
+async def test_elects_exactly_one_producer_per_stream_id() -> None:
+    store = create_in_memory_resumable_stream_store()
+    first = await store.acquire("a")
+    second = await store.acquire("a")
+    third = await store.acquire("a")
+    assert first == "producer"
+    assert second == "consumer"
+    assert third == "consumer"
+
+
+@pytest.mark.anyio
+async def test_stream_id_with_trailing_newline_is_invalid() -> None:
+    store = create_in_memory_resumable_stream_store()
+    with pytest.raises(ResumableStreamError) as exc:
+        await store.acquire("valid-id\n")
+    assert exc.value.code == "invalid-id"
+
+
+@pytest.mark.anyio
+async def test_post_finalize_acquire_is_consumer() -> None:
+    store = create_in_memory_resumable_stream_store()
+    assert await store.acquire("a") == "producer"
+    await store.finalize("a", "done")
+    assert await store.acquire("a") == "consumer"
+
+
+@pytest.mark.anyio
+async def test_isolates_streams_by_id() -> None:
+    store = create_in_memory_resumable_stream_store()
+    assert await store.acquire("a") == "producer"
+    assert await store.acquire("b") == "producer"
+
+
+@pytest.mark.anyio
+async def test_replays_buffered_entries_and_tails_until_finalize() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.append("a", _bytes("hello "))
+    await store.append("a", _bytes("world"))
+
+    signal = asyncio.Event()
+    collected: list[str] = []
+
+    async def reading() -> None:
+        async for entry in store.read("a", "", signal):
+            collected.append(_decode(entry.chunk))
+            if len(collected) == 3:
+                await store.finalize("a", "done")
+
+    task = asyncio.create_task(reading())
+    await store.append("a", _bytes("!"))
+    await task
+    assert collected == ["hello ", "world", "!"]
+
+
+@pytest.mark.anyio
+async def test_status_transitions_missing_streaming_done() -> None:
+    store = create_in_memory_resumable_stream_store()
+    assert await store.status("a") == "missing"
+    await store.acquire("a")
+    assert await store.status("a") == "streaming"
+    await store.finalize("a", "done")
+    assert await store.status("a") == "done"
+
+
+@pytest.mark.anyio
+async def test_status_reports_error_after_error_finalize() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.finalize("a", "error", "boom")
+    assert await store.status("a") == "error"
+
+
+@pytest.mark.anyio
+async def test_read_throws_after_error_finalize_after_draining() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.append("a", _bytes("partial"))
+    await store.finalize("a", "error", "boom")
+
+    signal = asyncio.Event()
+    seen: list[str] = []
+    with pytest.raises(Exception, match="boom"):
+        async for entry in store.read("a", "", signal):
+            seen.append(_decode(entry.chunk))
+    assert seen == ["partial"]
+
+
+@pytest.mark.anyio
+async def test_late_consumer_after_done_replays_everything() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.append("a", _bytes("a"))
+    await store.append("a", _bytes("b"))
+    await store.append("a", _bytes("c"))
+    await store.finalize("a", "done")
+
+    signal = asyncio.Event()
+    assert await _drain(store.read("a", "", signal)) == ["a", "b", "c"]
+
+
+@pytest.mark.anyio
+async def test_cursor_advances_and_skips_already_seen() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.append("a", _bytes("1"))
+    await store.append("a", _bytes("2"))
+    await store.append("a", _bytes("3"))
+    await store.finalize("a", "done")
+
+    signal = asyncio.Event()
+    seen: list[ResumableStreamEntry] = []
+    async for entry in store.read("a", "", signal):
+        seen.append(entry)
+    assert [_decode(s.chunk) for s in seen] == ["1", "2", "3"]
+
+    after_first = seen[0].cursor
+    assert await _drain(store.read("a", after_first, signal)) == ["2", "3"]
+
+
+@pytest.mark.anyio
+async def test_signal_set_terminates_read_without_raising() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    signal = asyncio.Event()
+    collected: list[str] = []
+
+    async def reading() -> None:
+        async for entry in store.read("a", "", signal):
+            collected.append(_decode(entry.chunk))
+
+    task = asyncio.create_task(reading())
+    await store.append("a", _bytes("x"))
+    await asyncio.sleep(0.01)
+    signal.set()
+    await task
+    assert collected == ["x"]
+
+
+@pytest.mark.anyio
+async def test_multiple_consumers_read_concurrently() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    signal = asyncio.Event()
+    a = asyncio.create_task(_drain(store.read("a", "", signal)))
+    b = asyncio.create_task(_drain(store.read("a", "", signal)))
+    await store.append("a", _bytes("x"))
+    await store.append("a", _bytes("y"))
+    await store.finalize("a", "done")
+    assert await a == ["x", "y"]
+    assert await b == ["x", "y"]
+
+
+@pytest.mark.anyio
+async def test_delete_ends_in_flight_reads_and_status_missing() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    signal = asyncio.Event()
+    reading = asyncio.create_task(_drain(store.read("a", "", signal)))
+    await store.append("a", _bytes("x"))
+    await asyncio.sleep(0)
+    await store.delete("a")
+    assert await reading == ["x"]
+    assert await store.status("a") == "missing"
+
+
+@pytest.mark.anyio
+async def test_expired_streams_evicted_on_next_access() -> None:
+    now = 1_000.0
+
+    def clock() -> float:
+        return now
+
+    store = create_in_memory_resumable_stream_store(default_ttl_ms=100, now=clock)
+    await store.acquire("a")
+    await store.append("a", _bytes("hi"))
+    assert await store.status("a") == "streaming"
+    now += 200
+    assert await store.status("a") == "missing"
+
+
+@pytest.mark.anyio
+async def test_appending_refreshes_ttl() -> None:
+    now = 1_000.0
+
+    def clock() -> float:
+        return now
+
+    store = create_in_memory_resumable_stream_store(default_ttl_ms=100, now=clock)
+    await store.acquire("a")
+    now += 80
+    await store.append("a", _bytes("x"))
+    now += 80
+    assert await store.status("a") == "streaming"
+
+
+@pytest.mark.anyio
+async def test_rejects_append_on_finalized_stream() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.finalize("a", "done")
+    with pytest.raises(ResumableStreamError, match="already finalized") as exc:
+        await store.append("a", _bytes("late"))
+    assert exc.value.code == "finalized"
+
+
+@pytest.mark.anyio
+async def test_rejects_append_on_missing_stream() -> None:
+    store = create_in_memory_resumable_stream_store()
+    with pytest.raises(Exception, match="Stream not found"):
+        await store.append("a", _bytes("x"))
+
+
+@pytest.mark.anyio
+async def test_finalize_is_idempotent() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.finalize("a", "done")
+    await store.finalize("a", "done")
+    assert await store.status("a") == "done"
+
+
+@pytest.mark.anyio
+async def test_rejects_append_when_chunk_exceeds_max_chunk_bytes() -> None:
+    store = create_in_memory_resumable_stream_store(max_chunk_bytes=4)
+    await store.acquire("a")
+    with pytest.raises(Exception, match="Chunk exceeds maxChunkBytes: 5"):
+        await store.append("a", _bytes("hello"))
+    await store.append("a", _bytes("ok"))
+    assert await store.status("a") == "streaming"
+
+
+@pytest.mark.anyio
+async def test_rejects_append_when_stream_reaches_max_entries() -> None:
+    store = create_in_memory_resumable_stream_store(max_entries_per_stream=2)
+    await store.acquire("a")
+    await store.append("a", _bytes("1"))
+    await store.append("a", _bytes("2"))
+    with pytest.raises(Exception, match="Stream exceeded maxEntriesPerStream: a"):
+        await store.append("a", _bytes("3"))
+
+
+@pytest.mark.anyio
+async def test_rejects_acquire_when_max_streams_exceeded() -> None:
+    store = create_in_memory_resumable_stream_store(max_streams=2)
+    await store.acquire("a")
+    await store.acquire("b")
+    with pytest.raises(Exception, match="maxStreams exceeded"):
+        await store.acquire("c")
+    assert await store.acquire("a") == "consumer"
+
+
+@pytest.mark.anyio
+async def test_gc_sweeper_evicts_expired_streams() -> None:
+    now = 1_000.0
+
+    def clock() -> float:
+        return now
+
+    store = create_in_memory_resumable_stream_store(
+        default_ttl_ms=100,
+        gc_interval_ms=50,
+        now=clock,
+    )
+    await store.acquire("a")
+    await store.append("a", _bytes("hi"))
+    now += 200
+    await asyncio.sleep(0.08)
+    assert await store.status("a") == "missing"
+    store.dispose()
+
+
+@pytest.mark.anyio
+async def test_dispose_clears_gc_and_is_noop_without_gc() -> None:
+    store = create_in_memory_resumable_stream_store(gc_interval_ms=50)
+    await store.acquire("a")
+    store.dispose()
+    store2 = create_in_memory_resumable_stream_store()
+    store2.dispose()