183 lines
6.2 KiB
Python
183 lines
6.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for the streaming SessionStore and BlobStore.
|
|
|
|
Covers:
|
|
|
|
* ``store`` / ``snapshot`` / ``drop`` lifecycle for the in-memory store
|
|
* ``hydrate`` with and without an explicit session id
|
|
* blob store insert / get / drop semantics
|
|
* thread-safety under concurrent writes (smoke)
|
|
* round-trip a LTX-2 continuation through snapshot + hydrate across a
|
|
session boundary (the "export and resume" flow the PR plan calls out)
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from fastvideo.api.schema import ContinuationState
|
|
from fastvideo.entrypoints.streaming.session_store import (
|
|
BlobStore,
|
|
InMemoryBlobStore,
|
|
InMemorySessionStore,
|
|
SessionStore,
|
|
)
|
|
from fastvideo.pipelines.basic.ltx2.continuation import (
|
|
LTX2_CONTINUATION_KIND,
|
|
LTX2ContinuationState,
|
|
)
|
|
|
|
|
|
class TestInMemoryBlobStore:
|
|
|
|
def test_is_blob_store(self):
|
|
assert isinstance(InMemoryBlobStore(), BlobStore)
|
|
|
|
def test_put_then_get_returns_same_bytes(self):
|
|
store = InMemoryBlobStore()
|
|
blob_id = store.put(b"hello")
|
|
assert store.get(blob_id) == b"hello"
|
|
|
|
def test_put_returns_distinct_ids(self):
|
|
store = InMemoryBlobStore()
|
|
id_a = store.put(b"a")
|
|
id_b = store.put(b"b")
|
|
assert id_a != id_b
|
|
|
|
def test_get_missing_raises_keyerror(self):
|
|
store = InMemoryBlobStore()
|
|
with pytest.raises(KeyError):
|
|
store.get("nonexistent")
|
|
|
|
def test_drop_removes_blob(self):
|
|
store = InMemoryBlobStore()
|
|
blob_id = store.put(b"payload")
|
|
store.drop(blob_id)
|
|
assert blob_id not in store
|
|
with pytest.raises(KeyError):
|
|
store.get(blob_id)
|
|
|
|
def test_drop_missing_is_noop(self):
|
|
store = InMemoryBlobStore()
|
|
store.drop("not-there") # no raise
|
|
|
|
def test_contains(self):
|
|
store = InMemoryBlobStore()
|
|
blob_id = store.put(b"x")
|
|
assert blob_id in store
|
|
assert "other" not in store
|
|
|
|
|
|
class TestInMemorySessionStore:
|
|
|
|
def test_is_session_store(self):
|
|
assert isinstance(InMemorySessionStore(), SessionStore)
|
|
|
|
def test_store_then_snapshot(self):
|
|
store = InMemorySessionStore()
|
|
state = ContinuationState(kind="ltx2.v1", payload={"x": 1})
|
|
store.store("sess-1", state)
|
|
assert store.snapshot("sess-1") is state
|
|
|
|
def test_snapshot_missing_returns_none(self):
|
|
store = InMemorySessionStore()
|
|
assert store.snapshot("missing") is None
|
|
|
|
def test_store_overwrites_prior_state(self):
|
|
store = InMemorySessionStore()
|
|
first = ContinuationState(kind="ltx2.v1", payload={"v": 1})
|
|
second = ContinuationState(kind="ltx2.v1", payload={"v": 2})
|
|
store.store("s", first)
|
|
store.store("s", second)
|
|
assert store.snapshot("s").payload["v"] == 2
|
|
|
|
def test_hydrate_assigns_new_session_id(self):
|
|
store = InMemorySessionStore()
|
|
state = ContinuationState(kind="ltx2.v1", payload={})
|
|
sid = store.hydrate(state)
|
|
assert sid
|
|
assert store.snapshot(sid) is state
|
|
|
|
def test_hydrate_with_explicit_session_id(self):
|
|
store = InMemorySessionStore()
|
|
state = ContinuationState(kind="ltx2.v1", payload={})
|
|
sid = store.hydrate(state, session_id="pinned-id")
|
|
assert sid == "pinned-id"
|
|
assert store.snapshot("pinned-id") is state
|
|
|
|
def test_drop_forgets_session(self):
|
|
store = InMemorySessionStore()
|
|
store.store("s", ContinuationState(kind="ltx2.v1", payload={}))
|
|
store.drop("s")
|
|
assert store.snapshot("s") is None
|
|
assert "s" not in store
|
|
|
|
def test_iter_yields_session_ids(self):
|
|
store = InMemorySessionStore()
|
|
store.store("a", ContinuationState(kind="ltx2.v1", payload={}))
|
|
store.store("b", ContinuationState(kind="ltx2.v1", payload={}))
|
|
assert sorted(store) == ["a", "b"]
|
|
|
|
def test_len(self):
|
|
store = InMemorySessionStore()
|
|
assert len(store) == 0
|
|
store.store("x", ContinuationState(kind="ltx2.v1", payload={}))
|
|
assert len(store) == 1
|
|
|
|
def test_concurrent_store_is_safe(self):
|
|
"""Smoke-check the lock: 200 parallel stores settle to 200 ids."""
|
|
store = InMemorySessionStore()
|
|
|
|
def write(i: int) -> None:
|
|
store.store(
|
|
f"s-{i}",
|
|
ContinuationState(kind="ltx2.v1", payload={"i": i}),
|
|
)
|
|
|
|
threads = [threading.Thread(target=write, args=(i, )) for i in range(200)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
assert len(store) == 200
|
|
|
|
|
|
class TestSnapshotHydrateRoundTrip:
|
|
"""Session boundary: snapshot + hydrate preserves the full LTX-2 state."""
|
|
|
|
def test_end_to_end_ltx2_session_migration(self):
|
|
blob_store = InMemoryBlobStore()
|
|
sessions = InMemorySessionStore()
|
|
|
|
typed = LTX2ContinuationState(
|
|
segment_index=4,
|
|
video_frames=[np.full((32, 32, 3), i * 5, dtype=np.uint8) for i in range(3)],
|
|
audio_latents=torch.randn(1, 4, 8, 32, dtype=torch.float32),
|
|
audio_sample_rate=24000,
|
|
audio_conditioning_num_frames=5,
|
|
video_position_offset_sec=0.25,
|
|
)
|
|
envelope = typed.to_continuation_state(blob_store=blob_store)
|
|
|
|
sessions.store("session-a", envelope)
|
|
snapshot = sessions.snapshot("session-a")
|
|
assert snapshot is not None
|
|
assert snapshot.kind == LTX2_CONTINUATION_KIND
|
|
|
|
# Simulate a migration: drop the first session, hydrate a new one
|
|
# from the snapshot, and reconstruct the typed state.
|
|
sessions.drop("session-a")
|
|
new_sid = sessions.hydrate(snapshot)
|
|
assert new_sid != "session-a"
|
|
rebuilt = sessions.snapshot(new_sid)
|
|
assert rebuilt is snapshot
|
|
|
|
restored = LTX2ContinuationState.from_continuation_state(rebuilt, blob_store=blob_store)
|
|
assert restored.segment_index == typed.segment_index
|
|
assert restored.audio_sample_rate == typed.audio_sample_rate
|
|
torch.testing.assert_close(restored.audio_latents, typed.audio_latents)
|
|
assert len(restored.video_frames) == len(typed.video_frames)
|