wehub-resource-sync
e30e75b5d4
Changesets / Create Version PR (push) Has been cancelled
Deploy Shadcn Registry / Deploy Production (push) Has been cancelled
Template Metrics / LOC + Bundle Size (push) Has been cancelled
Code Quality / Oxlint + Oxfmt (push) Has been cancelled
Code Quality / Template Sync (push) Has been cancelled
Code Quality / Build Changed Packages (push) Has been cancelled
Code Quality / Test Changed Packages (push) Has been cancelled
Deploy Expo Example / Deploy Production (push) Has been cancelled
Deploy Ink Example / Deploy Production (push) Has been cancelled
Python Tests / pytest (assistant-stream, 3.10) (push) Has been cancelled
Python Tests / pytest (assistant-stream, 3.12) (push) Has been cancelled
Python Tests / pytest (assistant-ui-sync-server-api, 3.10) (push) Has been cancelled
Python Tests / pytest (assistant-ui-sync-server-api, 3.12) (push) Has been cancelled
309 lines
9.5 KiB
Python
309 lines
9.5 KiB
Python
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()
|