Coverage for tests/io/test_stream.py: 99%
211 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-18 05:24 +0000
1"""Streaming API tests: EventStreamer (evutils.io.stream) and the pipeline
2generators in evutils.chunking (stream_delta_t / stream_n_events /
3stream_skip_to_time / stream_async / stream_paced_playback).
5EventStreamer is exercised over a real encoder-written file; the pipeline
6generators are driven with hand-built EventArray/TriggerArray chunks so every
7branch (trigger vs no-trigger, boundary slicing, leftover flush, empty chunks,
8exception propagation) is covered without needing a device.
9"""
10import time
12import numpy as np
13import pytest
14from typing import Any
16from evutils.types import EventArray, TriggerArray, Event_dtype
17from evutils.chunking import (
18 stream_delta_t,
19 stream_n_events,
20 stream_skip_to_time,
21 stream_async,
22 stream_paced_playback,
23)
26# --------------------------------------------------------------------------- #
27# helpers
28# --------------------------------------------------------------------------- #
30def ev(ts, xs=None, ys=None, ps=None) -> EventArray:
31 ts = list(ts)
32 n = len(ts)
33 return EventArray(
34 t=ts,
35 x=xs if xs is not None else [0] * n,
36 y=ys if ys is not None else [0] * n,
37 p=ps if ps is not None else [0] * n,
38 )
41def tr(ts, ids=None, ps=None) -> TriggerArray:
42 ts = list(ts)
43 n = len(ts)
44 return TriggerArray(t=ts, p=ps if ps is not None else [0] * n, id=ids if ids is not None else [0] * n)
47def cat_t(chunks) -> list:
48 """Flatten a list of chunks (or (ev, tr) tuples) into their event timestamps."""
49 out = []
50 for c in chunks:
51 e = c[0] if isinstance(c, tuple) else c
52 out.extend(e.t.tolist())
53 return out
56def cat_tr_t(chunks) -> list:
57 out = []
58 for c in chunks:
59 out.extend(c[1].t.tolist())
60 return out
63# --------------------------------------------------------------------------- #
64# stream_delta_t
65# --------------------------------------------------------------------------- #
67def test_stream_delta_t_regroups_across_chunk_boundaries():
68 raw = [ev([0, 10, 20]), ev([30, 35, 55]), ev([60, 90, 95])]
69 out = list(stream_delta_t(iter(raw), delta_t=30))
70 # A trigger-less stream yields bare EventArrays (the output mirrors the
71 # input shape). Every input event survives, in order.
72 assert all(not isinstance(w, tuple) for w in out)
73 assert cat_t(out) == [0, 10, 20, 30, 35, 55, 60, 90, 95]
74 # full (non-final) windows respect the delta_t span
75 for e in out[:-1]:
76 if len(e) > 1:
77 assert int(e.t[-1]) - int(e.t[0]) < 30
80def test_stream_delta_t_with_triggers():
81 raw = [
82 (ev([0, 10]), tr([5], ids=[1])),
83 (ev([40, 50]), tr([45], ids=[2])),
84 ]
85 out = list(stream_delta_t(iter(raw), delta_t=30))
86 assert all(isinstance(c, tuple) for c in out)
87 assert cat_t(out) == [0, 10, 40, 50]
88 assert cat_tr_t(out) == [5, 45]
91def test_stream_delta_t_empty_stream():
92 assert list(stream_delta_t(iter([]), delta_t=30)) == []
95def test_stream_delta_t_skips_empty_leading_chunks():
96 raw = [ev([]), ev([1, 2, 3])]
97 out = list(stream_delta_t(iter(raw), delta_t=100))
98 assert cat_t(out) == [1, 2, 3]
101# --------------------------------------------------------------------------- #
102# stream_n_events
103# --------------------------------------------------------------------------- #
105def test_stream_n_events_fixed_size_with_remainder():
106 raw = [ev(range(0, 5)), ev(range(5, 10)), ev(range(10, 13))]
107 out = list(stream_n_events(iter(raw), n_events=4))
108 # trigger-less input -> bare EventArray windows (output mirrors input shape)
109 assert all(not isinstance(w, tuple) for w in out)
110 assert [len(w) for w in out] == [4, 4, 4, 1] # 13 = 3*4 + 1
111 assert cat_t(out) == list(range(13))
114def test_stream_n_events_with_triggers_exact_boundary():
115 # total events an exact multiple of n_events -> exact-boundary trigger branch
116 raw = [(ev([0, 1, 2, 3]), tr([1, 3], ids=[1, 2]))]
117 out = list(stream_n_events(iter(raw), n_events=4))
118 assert len(out) == 1
119 e, t = out[0]
120 assert e.t.tolist() == [0, 1, 2, 3]
121 assert t.t.tolist() == [1, 3]
124def test_stream_n_events_with_triggers_searchsorted_branch():
125 # more events than n_events in the buffer -> searchsorted trigger split
126 raw = [(ev([0, 1, 2, 3, 4, 5]), tr([0, 2, 5], ids=[1, 2, 3]))]
127 out = list(stream_n_events(iter(raw), n_events=4))
128 assert cat_t(out) == [0, 1, 2, 3, 4, 5]
129 # triggers before t of the 4th event go with the first window
130 assert cat_tr_t(out) == [0, 2, 5]
133def test_stream_n_events_empty_stream():
134 assert list(stream_n_events(iter([]), n_events=4)) == []
137# --------------------------------------------------------------------------- #
138# stream_skip_to_time
139# --------------------------------------------------------------------------- #
141def test_stream_skip_to_time_drops_whole_and_partial_chunks():
142 raw = [ev([0, 10, 20]), ev([30, 40, 50])]
143 out = list(stream_skip_to_time(iter(raw), start_ts=25))
144 assert cat_t(out) == [30, 40, 50] # first chunk dropped whole
147def test_stream_skip_to_time_partial_slice_within_chunk():
148 raw = [ev([0, 10, 20]), ev([30, 40])]
149 out = list(stream_skip_to_time(iter(raw), start_ts=15))
150 assert cat_t(out) == [20, 30, 40] # first chunk sliced at 15
153def test_stream_skip_to_time_with_triggers():
154 raw = [
155 (ev([0, 10, 20]), tr([5, 18], ids=[1, 2])),
156 (ev([30, 40]), tr([35], ids=[3])),
157 ]
158 out = list(stream_skip_to_time(iter(raw), start_ts=15))
159 assert cat_t(out) == [20, 30, 40]
160 assert cat_tr_t(out) == [18, 35] # trigger at 5 dropped
163def test_stream_skip_to_time_never_reached_yields_nothing():
164 raw = [ev([0, 10]), ev([20, 30])]
165 out = list(stream_skip_to_time(iter(raw), start_ts=1000))
166 assert out == []
169# --------------------------------------------------------------------------- #
170# stream_async
171# --------------------------------------------------------------------------- #
173def test_stream_async_preserves_order():
174 raw = [ev([0, 1]), ev([2, 3]), ev([4])]
175 out = list(stream_async(iter(raw)))
176 assert cat_t(out) == [0, 1, 2, 3, 4]
179def test_stream_async_with_triggers_and_none():
180 raw = [(ev([0, 1]), tr([0], ids=[1])), (ev([2]), None)]
181 out = list(stream_async(iter(raw)))
182 assert cat_t(out) == [0, 1, 2]
183 assert out[0][1].t.tolist() == [0]
184 assert out[1][1] is None
187def test_stream_async_copies_chunks():
188 """Chunks are copied before crossing the thread, so mutating the source
189 afterwards must not change what was yielded."""
190 src = ev([7, 8])
191 out = list(stream_async(iter([src])))
192 src.t[0] = 999
193 assert out[0].t.tolist() == [7, 8]
196def test_stream_async_propagates_exceptions():
197 def boom():
198 yield ev([0])
199 raise ValueError("upstream failed")
201 with pytest.raises(ValueError, match="upstream failed"):
202 list(stream_async(boom()))
205# --------------------------------------------------------------------------- #
206# stream_paced_playback
207# --------------------------------------------------------------------------- #
209def test_stream_paced_playback_preserves_data_fast():
210 raw = [ev([0, 100]), ev([200, 300])]
211 out = list(stream_paced_playback(iter(raw), playback_speed=1e6)) # ~no sleep
212 assert cat_t(out) == [0, 100, 200, 300]
215def test_stream_paced_playback_empty_chunk_passthrough():
216 raw = [ev([]), ev([0, 10])]
217 out = list(stream_paced_playback(iter(raw), playback_speed=1e6))
218 assert cat_t(out) == [0, 10]
221def test_stream_paced_playback_actually_waits():
222 # 20 ms of stream at 2x speed -> ~10 ms of real wait (the sleep branch).
223 raw = [ev([0]), ev([20_000])]
224 start = time.perf_counter()
225 out = list(stream_paced_playback(iter(raw), playback_speed=2.0))
226 elapsed = time.perf_counter() - start
227 assert cat_t(out) == [0, 20_000]
228 assert elapsed >= 0.005 # some real delay occurred
231# --------------------------------------------------------------------------- #
232# EventStreamer (integration over a written file)
233# --------------------------------------------------------------------------- #
235def _make_events(n=500):
236 a = np.zeros(n, dtype=Event_dtype)
237 a['t'] = np.arange(n, dtype=np.int64) * 100
238 a['x'] = np.arange(n) % 1280
239 a['y'] = np.arange(n) % 720
240 a['p'] = np.arange(n) % 2
241 return a
244def test_event_streamer_yields_all_events(tmp_path):
245 from evutils.io import EventWriter, EventStreamer
246 events = _make_events()
247 p = tmp_path / "stream.raw"
248 with EventWriter(p, format="evt3") as w:
249 w.write(events)
251 # EventStreamer yields chunks that alias a reused parser buffer, so retain a
252 # copy per iteration (the documented low-level contract; stream_async does
253 # the same copy).
254 ts = []
255 n = 0
256 for c in EventStreamer(p):
257 ts.append(c.t.copy())
258 n += 1
259 assert n > 0
260 got = np.concatenate(ts)
261 assert np.array_equal(got, events['t'])
264def test_event_streamer_pipes_into_stream_n_events(tmp_path):
265 from evutils.io import EventWriter, EventStreamer
266 events = _make_events()
267 p = tmp_path / "stream.raw"
268 with EventWriter(p, format="evt3") as w:
269 w.write(events)
271 # stream_n_events copies into its accumulator, so piping the aliasing
272 # EventStreamer through it is safe. Trigger-less input -> bare windows.
273 windows = list(stream_n_events(EventStreamer(p), n_events=128))
274 assert [len(w) for w in windows[:-1]] == [128] * (len(windows) - 1)
275 got = np.concatenate([w.t for w in windows])
276 assert np.array_equal(got, events['t'])
279class _FakeDecoder:
280 """Minimal decoder that yields one plain (non-tuple) event chunk then stops,
281 used to drive EventStreamer with an explicit decoder_cls and to exercise the
282 ext_trigger path when the decoder returns a bare EventArray."""
283 def __init__(self, source, read_external_triggers=False, **kwargs):
284 self._chunks = iter([ev([1, 2, 3]), ev([])])
286 def init(self):
287 pass
289 def read_chunk(self):
290 return next(self._chunks)
293def test_event_streamer_explicit_decoder_cls_and_bare_chunk(tmp_path):
294 from evutils.io import EventStreamer
295 p = tmp_path / "any.raw"
296 p.write_bytes(b"% end\n") # make_source needs a real source; content unused
298 # ext_trigger=True but the decoder returns a bare EventArray -> streamer must
299 # pair it with an empty TriggerArray.
300 out = list(EventStreamer(p, ext_trigger=True, decoder_cls=_FakeDecoder))
301 assert len(out) == 1
302 e, t = out[0]
303 assert e.t.tolist() == [1, 2, 3]
304 assert len(t) == 0
306 # no-trigger path with the same explicit decoder_cls
307 out2 = list(EventStreamer(p, decoder_cls=_FakeDecoder))
308 assert [c.t.tolist() for c in out2] == [[1, 2, 3]]
311def test_event_streamer_with_triggers(tmp_path):
312 from evutils.io import EventWriter, EventStreamer
313 events = _make_events()
314 p = tmp_path / "stream.raw"
315 with EventWriter(p, format="evt3") as w:
316 w.write(events)
318 ts, total_tr, all_tuples = [], 0, True
319 for c in EventStreamer(p, ext_trigger=True):
320 all_tuples = all_tuples and isinstance(c, tuple)
321 ts.append(c[0].t.copy()) # copy: chunk aliases the reused parser buffer
322 total_tr += len(c[1])
323 assert all_tuples
324 got = np.concatenate(ts)
325 assert np.array_equal(got, events['t'])
326 assert total_tr == 0 # no trigger words in the file
329def test_stream_tuple_unwrapping_when_ext_trigger_false(tmp_path: Any) -> None:
330 from evutils.io.stream import EventStreamer
331 from evutils.io.decoders import EventDecoder
332 from evutils.types import EventArray, TriggerArray
334 class MockDecoder(EventDecoder):
335 def __init__(self, source: Any, **kwargs: Any) -> None:
336 self.source = source
337 self.read_external_triggers = False
338 self.called = False
340 def init(self) -> None:
341 pass
343 def reset(self) -> None:
344 self.called = False
346 def read_chunk(self, delta_t: int | None = None, n_events_hint: int | None = None) -> Any:
347 if self.called:
348 return EventArray.empty()
349 self.called = True
350 ev = EventArray(t=[1], x=[2], y=[3], p=[0])
351 tr = TriggerArray(t=[1], p=[1], id=[0])
352 return ev, tr # Unconditionally returns tuple!
354 def is_eof(self) -> bool:
355 return self.called
357 p = tmp_path / "mock.raw"
358 p.touch()
359 streamer = EventStreamer(p, ext_trigger=False, decoder_cls=MockDecoder)
360 chunks = list(streamer)
361 assert len(chunks) == 1
362 assert isinstance(chunks[0], EventArray)
363 assert len(chunks[0]) == 1