Coverage for tests/io/test_stream.py: 100%
180 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-15 09:33 +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
15from evutils.types import EventArray, TriggerArray, Event_dtype
16from evutils.chunking import (
17 stream_delta_t,
18 stream_n_events,
19 stream_skip_to_time,
20 stream_async,
21 stream_paced_playback,
22)
25# --------------------------------------------------------------------------- #
26# helpers
27# --------------------------------------------------------------------------- #
29def ev(ts, xs=None, ys=None, ps=None) -> EventArray:
30 ts = list(ts)
31 n = len(ts)
32 return EventArray(
33 t=ts,
34 x=xs if xs is not None else [0] * n,
35 y=ys if ys is not None else [0] * n,
36 p=ps if ps is not None else [0] * n,
37 )
40def tr(ts, ids=None, ps=None) -> TriggerArray:
41 ts = list(ts)
42 n = len(ts)
43 return TriggerArray(t=ts, p=ps if ps is not None else [0] * n, id=ids if ids is not None else [0] * n)
46def cat_t(chunks) -> list:
47 """Flatten a list of chunks (or (ev, tr) tuples) into their event timestamps."""
48 out = []
49 for c in chunks:
50 e = c[0] if isinstance(c, tuple) else c
51 out.extend(e.t.tolist())
52 return out
55def cat_tr_t(chunks) -> list:
56 out = []
57 for c in chunks:
58 out.extend(c[1].t.tolist())
59 return out
62# --------------------------------------------------------------------------- #
63# stream_delta_t
64# --------------------------------------------------------------------------- #
66def test_stream_delta_t_regroups_across_chunk_boundaries():
67 raw = [ev([0, 10, 20]), ev([30, 35, 55]), ev([60, 90, 95])]
68 out = list(stream_delta_t(iter(raw), delta_t=30))
69 # stream_delta_t always yields (ev, tr) tuples (the accumulator always has a
70 # trigger buffer). Every input event survives, in order.
71 assert cat_t(out) == [0, 10, 20, 30, 35, 55, 60, 90, 95]
72 # full (non-final) windows respect the delta_t span
73 for w in out[:-1]:
74 e = w[0]
75 if len(e) > 1:
76 assert int(e.t[-1]) - int(e.t[0]) < 30
79def test_stream_delta_t_with_triggers():
80 raw = [
81 (ev([0, 10]), tr([5], ids=[1])),
82 (ev([40, 50]), tr([45], ids=[2])),
83 ]
84 out = list(stream_delta_t(iter(raw), delta_t=30))
85 assert all(isinstance(c, tuple) for c in out)
86 assert cat_t(out) == [0, 10, 40, 50]
87 assert cat_tr_t(out) == [5, 45]
90def test_stream_delta_t_empty_stream():
91 assert list(stream_delta_t(iter([]), delta_t=30)) == []
94def test_stream_delta_t_skips_empty_leading_chunks():
95 raw = [ev([]), ev([1, 2, 3])]
96 out = list(stream_delta_t(iter(raw), delta_t=100))
97 assert cat_t(out) == [1, 2, 3]
100# --------------------------------------------------------------------------- #
101# stream_n_events
102# --------------------------------------------------------------------------- #
104def test_stream_n_events_fixed_size_with_remainder():
105 raw = [ev(range(0, 5)), ev(range(5, 10)), ev(range(10, 13))]
106 out = list(stream_n_events(iter(raw), n_events=4))
107 # always (ev, tr) tuples; window sizes counted on the event component
108 assert [len(w[0]) for w in out] == [4, 4, 4, 1] # 13 = 3*4 + 1
109 assert cat_t(out) == list(range(13))
112def test_stream_n_events_with_triggers_exact_boundary():
113 # total events an exact multiple of n_events -> exact-boundary trigger branch
114 raw = [(ev([0, 1, 2, 3]), tr([1, 3], ids=[1, 2]))]
115 out = list(stream_n_events(iter(raw), n_events=4))
116 assert len(out) == 1
117 e, t = out[0]
118 assert e.t.tolist() == [0, 1, 2, 3]
119 assert t.t.tolist() == [1, 3]
122def test_stream_n_events_with_triggers_searchsorted_branch():
123 # more events than n_events in the buffer -> searchsorted trigger split
124 raw = [(ev([0, 1, 2, 3, 4, 5]), tr([0, 2, 5], ids=[1, 2, 3]))]
125 out = list(stream_n_events(iter(raw), n_events=4))
126 assert cat_t(out) == [0, 1, 2, 3, 4, 5]
127 # triggers before t of the 4th event go with the first window
128 assert cat_tr_t(out) == [0, 2, 5]
131def test_stream_n_events_empty_stream():
132 assert list(stream_n_events(iter([]), n_events=4)) == []
135# --------------------------------------------------------------------------- #
136# stream_skip_to_time
137# --------------------------------------------------------------------------- #
139def test_stream_skip_to_time_drops_whole_and_partial_chunks():
140 raw = [ev([0, 10, 20]), ev([30, 40, 50])]
141 out = list(stream_skip_to_time(iter(raw), start_ts=25))
142 assert cat_t(out) == [30, 40, 50] # first chunk dropped whole
145def test_stream_skip_to_time_partial_slice_within_chunk():
146 raw = [ev([0, 10, 20]), ev([30, 40])]
147 out = list(stream_skip_to_time(iter(raw), start_ts=15))
148 assert cat_t(out) == [20, 30, 40] # first chunk sliced at 15
151def test_stream_skip_to_time_with_triggers():
152 raw = [
153 (ev([0, 10, 20]), tr([5, 18], ids=[1, 2])),
154 (ev([30, 40]), tr([35], ids=[3])),
155 ]
156 out = list(stream_skip_to_time(iter(raw), start_ts=15))
157 assert cat_t(out) == [20, 30, 40]
158 assert cat_tr_t(out) == [18, 35] # trigger at 5 dropped
161def test_stream_skip_to_time_never_reached_yields_nothing():
162 raw = [ev([0, 10]), ev([20, 30])]
163 out = list(stream_skip_to_time(iter(raw), start_ts=1000))
164 assert out == []
167# --------------------------------------------------------------------------- #
168# stream_async
169# --------------------------------------------------------------------------- #
171def test_stream_async_preserves_order():
172 raw = [ev([0, 1]), ev([2, 3]), ev([4])]
173 out = list(stream_async(iter(raw)))
174 assert cat_t(out) == [0, 1, 2, 3, 4]
177def test_stream_async_with_triggers_and_none():
178 raw = [(ev([0, 1]), tr([0], ids=[1])), (ev([2]), None)]
179 out = list(stream_async(iter(raw)))
180 assert cat_t(out) == [0, 1, 2]
181 assert out[0][1].t.tolist() == [0]
182 assert out[1][1] is None
185def test_stream_async_copies_chunks():
186 """Chunks are copied before crossing the thread, so mutating the source
187 afterwards must not change what was yielded."""
188 src = ev([7, 8])
189 out = list(stream_async(iter([src])))
190 src.t[0] = 999
191 assert out[0].t.tolist() == [7, 8]
194def test_stream_async_propagates_exceptions():
195 def boom():
196 yield ev([0])
197 raise ValueError("upstream failed")
199 with pytest.raises(ValueError, match="upstream failed"):
200 list(stream_async(boom()))
203# --------------------------------------------------------------------------- #
204# stream_paced_playback
205# --------------------------------------------------------------------------- #
207def test_stream_paced_playback_preserves_data_fast():
208 raw = [ev([0, 100]), ev([200, 300])]
209 out = list(stream_paced_playback(iter(raw), playback_speed=1e6)) # ~no sleep
210 assert cat_t(out) == [0, 100, 200, 300]
213def test_stream_paced_playback_empty_chunk_passthrough():
214 raw = [ev([]), ev([0, 10])]
215 out = list(stream_paced_playback(iter(raw), playback_speed=1e6))
216 assert cat_t(out) == [0, 10]
219def test_stream_paced_playback_actually_waits():
220 # 20 ms of stream at 2x speed -> ~10 ms of real wait (the sleep branch).
221 raw = [ev([0]), ev([20_000])]
222 start = time.perf_counter()
223 out = list(stream_paced_playback(iter(raw), playback_speed=2.0))
224 elapsed = time.perf_counter() - start
225 assert cat_t(out) == [0, 20_000]
226 assert elapsed >= 0.005 # some real delay occurred
229# --------------------------------------------------------------------------- #
230# EventStreamer (integration over a written file)
231# --------------------------------------------------------------------------- #
233def _make_events(n=500):
234 a = np.zeros(n, dtype=Event_dtype)
235 a['t'] = np.arange(n, dtype=np.int64) * 100
236 a['x'] = np.arange(n) % 1280
237 a['y'] = np.arange(n) % 720
238 a['p'] = np.arange(n) % 2
239 return a
242def test_event_streamer_yields_all_events(tmp_path):
243 from evutils.io import EventWriter, EventStreamer
244 events = _make_events()
245 p = tmp_path / "stream.raw"
246 with EventWriter(p, format="evt3") as w:
247 w.write(events)
249 # EventStreamer yields chunks that alias a reused parser buffer, so retain a
250 # copy per iteration (the documented low-level contract; stream_async does
251 # the same copy).
252 ts = []
253 n = 0
254 for c in EventStreamer(p):
255 ts.append(c.t.copy())
256 n += 1
257 assert n > 0
258 got = np.concatenate(ts)
259 assert np.array_equal(got, events['t'])
262def test_event_streamer_pipes_into_stream_n_events(tmp_path):
263 from evutils.io import EventWriter, EventStreamer
264 events = _make_events()
265 p = tmp_path / "stream.raw"
266 with EventWriter(p, format="evt3") as w:
267 w.write(events)
269 # stream_n_events copies into its accumulator, so piping the aliasing
270 # EventStreamer through it is safe. Output windows are (ev, tr) tuples.
271 windows = list(stream_n_events(EventStreamer(p), n_events=128))
272 assert [len(w[0]) for w in windows[:-1]] == [128] * (len(windows) - 1)
273 got = np.concatenate([w[0].t for w in windows])
274 assert np.array_equal(got, events['t'])
277class _FakeDecoder:
278 """Minimal decoder that yields one plain (non-tuple) event chunk then stops,
279 used to drive EventStreamer with an explicit decoder_cls and to exercise the
280 ext_trigger path when the decoder returns a bare EventArray."""
281 def __init__(self, source, read_external_triggers=False, **kwargs):
282 self._chunks = iter([ev([1, 2, 3]), ev([])])
284 def init(self):
285 pass
287 def read_chunk(self):
288 return next(self._chunks)
291def test_event_streamer_explicit_decoder_cls_and_bare_chunk(tmp_path):
292 from evutils.io import EventStreamer
293 p = tmp_path / "any.raw"
294 p.write_bytes(b"% end\n") # make_source needs a real source; content unused
296 # ext_trigger=True but the decoder returns a bare EventArray -> streamer must
297 # pair it with an empty TriggerArray.
298 out = list(EventStreamer(p, ext_trigger=True, decoder_cls=_FakeDecoder))
299 assert len(out) == 1
300 e, t = out[0]
301 assert e.t.tolist() == [1, 2, 3]
302 assert len(t) == 0
304 # no-trigger path with the same explicit decoder_cls
305 out2 = list(EventStreamer(p, decoder_cls=_FakeDecoder))
306 assert [c.t.tolist() for c in out2] == [[1, 2, 3]]
309def test_event_streamer_with_triggers(tmp_path):
310 from evutils.io import EventWriter, EventStreamer
311 events = _make_events()
312 p = tmp_path / "stream.raw"
313 with EventWriter(p, format="evt3") as w:
314 w.write(events)
316 ts, total_tr, all_tuples = [], 0, True
317 for c in EventStreamer(p, ext_trigger=True):
318 all_tuples = all_tuples and isinstance(c, tuple)
319 ts.append(c[0].t.copy()) # copy: chunk aliases the reused parser buffer
320 total_tr += len(c[1])
321 assert all_tuples
322 got = np.concatenate(ts)
323 assert np.array_equal(got, events['t'])
324 assert total_tr == 0 # no trigger words in the file