bt2: refactor test_message_iterator
[babeltrace.git] / tests / bindings / python / bt2 / test_message_iterator.py
CommitLineData
d2d857a8
MJ
1#
2# Copyright (C) 2019 EfficiOS Inc.
3#
4# This program is free software; you can redistribute it and/or
5# modify it under the terms of the GNU General Public License
6# as published by the Free Software Foundation; only version 2
7# of the License.
8#
9# This program is distributed in the hope that it will be useful,
10# but WITHOUT ANY WARRANTY; without even the implied warranty of
11# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12# GNU General Public License for more details.
13#
14# You should have received a copy of the GNU General Public License
15# along with this program; if not, write to the Free Software
16# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
17#
18
811644b8 19import unittest
811644b8 20import bt2
8e97c333 21import sys
6c373cc9 22from utils import TestOutputPortMessageIterator
14503fb1 23from bt2 import port as bt2_port
811644b8
PP
24
25
0a6d7302
SM
26class SimpleSink(bt2._UserSinkComponent):
27 # Straightforward sink that creates one input port (`in`) and consumes from
28 # it.
811644b8 29
0a6d7302
SM
30 def __init__(self, params, obj):
31 self._add_input_port('in')
811644b8 32
0a6d7302
SM
33 def _user_consume(self):
34 next(self._msg_iter)
811644b8 35
0a6d7302
SM
36 def _user_graph_is_configured(self):
37 self._msg_iter = self._create_input_port_message_iterator(
38 self._input_ports['in']
39 )
ca02df0a 40
ca02df0a 41
0a6d7302
SM
42def _create_graph(src_comp_cls, sink_comp_cls, flt_comp_cls=None):
43 graph = bt2.Graph()
ca02df0a 44
0a6d7302
SM
45 src_comp = graph.add_component(src_comp_cls, 'src')
46 sink_comp = graph.add_component(sink_comp_cls, 'sink')
ca02df0a 47
0a6d7302
SM
48 if flt_comp_cls is not None:
49 flt_comp = graph.add_component(flt_comp_cls, 'flt')
50 graph.connect_ports(src_comp.output_ports['out'], flt_comp.input_ports['in'])
51 graph.connect_ports(flt_comp.output_ports['out'], sink_comp.input_ports['in'])
52 else:
53 graph.connect_ports(src_comp.output_ports['out'], sink_comp.input_ports['in'])
811644b8 54
0a6d7302
SM
55 return graph
56
57
58class UserMessageIteratorTestCase(unittest.TestCase):
811644b8 59 def test_init(self):
c5f330cd
SM
60 the_output_port_from_source = None
61 the_output_port_from_iter = None
62
5602ef81 63 class MyIter(bt2._UserMessageIterator):
c5f330cd 64 def __init__(self, self_port_output):
811644b8 65 nonlocal initialized
c5f330cd 66 nonlocal the_output_port_from_iter
811644b8 67 initialized = True
c5f330cd 68 the_output_port_from_iter = self_port_output
811644b8 69
cfbd7cf3 70 class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
66964f3f 71 def __init__(self, params, obj):
c5f330cd 72 nonlocal the_output_port_from_source
2e00bc76 73 the_output_port_from_source = self._add_output_port('out', 'user data')
811644b8
PP
74
75 initialized = False
0a6d7302 76 graph = _create_graph(MySource, SimpleSink)
c5f330cd 77 graph.run()
811644b8 78 self.assertTrue(initialized)
cfbd7cf3
FD
79 self.assertEqual(
80 the_output_port_from_source.addr, the_output_port_from_iter.addr
81 )
2e00bc76 82 self.assertEqual(the_output_port_from_iter.user_data, 'user data')
811644b8 83
ca02df0a
PP
84 def test_create_from_message_iterator(self):
85 class MySourceIter(bt2._UserMessageIterator):
86 def __init__(self, self_port_output):
87 nonlocal src_iter_initialized
88 src_iter_initialized = True
89
90 class MySource(bt2._UserSourceComponent, message_iterator_class=MySourceIter):
66964f3f 91 def __init__(self, params, obj):
ca02df0a
PP
92 self._add_output_port('out')
93
94 class MyFilterIter(bt2._UserMessageIterator):
95 def __init__(self, self_port_output):
96 nonlocal flt_iter_initialized
97 flt_iter_initialized = True
98 self._up_iter = self._create_input_port_message_iterator(
99 self._component._input_ports['in']
100 )
101
102 def __next__(self):
103 return next(self._up_iter)
104
105 class MyFilter(bt2._UserFilterComponent, message_iterator_class=MyFilterIter):
66964f3f 106 def __init__(self, params, obj):
ca02df0a
PP
107 self._add_input_port('in')
108 self._add_output_port('out')
109
110 src_iter_initialized = False
111 flt_iter_initialized = False
0a6d7302 112 graph = _create_graph(MySource, SimpleSink, MyFilter)
ca02df0a
PP
113 graph.run()
114 self.assertTrue(src_iter_initialized)
115 self.assertTrue(flt_iter_initialized)
116
e803df70
SM
117 def test_create_user_error(self):
118 # This tests both error handling by
119 # _UserSinkComponent._create_input_port_message_iterator
120 # and _UserMessageIterator._create_input_port_message_iterator, as they
121 # are both used in the graph.
122 class MySourceIter(bt2._UserMessageIterator):
123 def __init__(self, self_port_output):
124 raise ValueError('Very bad error')
125
126 class MySource(bt2._UserSourceComponent, message_iterator_class=MySourceIter):
127 def __init__(self, params, obj):
128 self._add_output_port('out')
129
130 class MyFilterIter(bt2._UserMessageIterator):
131 def __init__(self, self_port_output):
132 # This is expected to raise because of the error in
133 # MySourceIter.__init__.
134 self._create_input_port_message_iterator(
135 self._component._input_ports['in']
136 )
137
138 class MyFilter(bt2._UserFilterComponent, message_iterator_class=MyFilterIter):
139 def __init__(self, params, obj):
140 self._add_input_port('in')
141 self._add_output_port('out')
142
0a6d7302 143 graph = _create_graph(MySource, SimpleSink, MyFilter)
e803df70
SM
144
145 with self.assertRaises(bt2._Error) as ctx:
146 graph.run()
147
148 exc = ctx.exception
149 cause = exc[0]
150
151 self.assertIsInstance(cause, bt2._MessageIteratorErrorCause)
152 self.assertEqual(cause.component_name, 'src')
153 self.assertEqual(cause.component_output_port_name, 'out')
154 self.assertIn('ValueError: Very bad error', cause.message)
155
811644b8 156 def test_finalize(self):
5602ef81 157 class MyIter(bt2._UserMessageIterator):
6a91742b 158 def _user_finalize(self):
811644b8
PP
159 nonlocal finalized
160 finalized = True
161
cfbd7cf3 162 class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
66964f3f 163 def __init__(self, params, obj):
811644b8
PP
164 self._add_output_port('out')
165
166 finalized = False
0a6d7302 167 graph = _create_graph(MySource, SimpleSink)
c5f330cd 168 graph.run()
811644b8
PP
169 del graph
170 self.assertTrue(finalized)
171
172 def test_component(self):
5602ef81 173 class MyIter(bt2._UserMessageIterator):
c5f330cd 174 def __init__(self, self_port_output):
811644b8
PP
175 nonlocal salut
176 salut = self._component._salut
177
cfbd7cf3 178 class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
66964f3f 179 def __init__(self, params, obj):
811644b8
PP
180 self._add_output_port('out')
181 self._salut = 23
182
183 salut = None
0a6d7302 184 graph = _create_graph(MySource, SimpleSink)
c5f330cd 185 graph.run()
811644b8
PP
186 self.assertEqual(salut, 23)
187
14503fb1
SM
188 def test_port(self):
189 class MyIter(bt2._UserMessageIterator):
190 def __init__(self_iter, self_port_output):
191 nonlocal called
192 called = True
193 port = self_iter._port
194 self.assertIs(type(self_port_output), bt2_port._UserComponentOutputPort)
195 self.assertIs(type(port), bt2_port._UserComponentOutputPort)
196 self.assertEqual(self_port_output.addr, port.addr)
197
198 class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
199 def __init__(self, params, obj):
200 self._add_output_port('out')
201
202 called = False
0a6d7302 203 graph = _create_graph(MySource, SimpleSink)
14503fb1
SM
204 graph.run()
205 self.assertTrue(called)
206
811644b8 207 def test_addr(self):
5602ef81 208 class MyIter(bt2._UserMessageIterator):
c5f330cd 209 def __init__(self, self_port_output):
811644b8
PP
210 nonlocal addr
211 addr = self.addr
212
cfbd7cf3 213 class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
66964f3f 214 def __init__(self, params, obj):
811644b8
PP
215 self._add_output_port('out')
216
217 addr = None
0a6d7302 218 graph = _create_graph(MySource, SimpleSink)
c5f330cd 219 graph.run()
811644b8
PP
220 self.assertIsNotNone(addr)
221 self.assertNotEqual(addr, 0)
222
d79a8353
SM
223 # Test that messages returned by _UserMessageIterator.__next__ remain valid
224 # and can be re-used.
225 def test_reuse_message(self):
226 class MyIter(bt2._UserMessageIterator):
227 def __init__(self, port):
228 tc, sc, ec = port.user_data
229 trace = tc()
230 stream = trace.create_stream(sc)
231 packet = stream.create_packet()
232
233 # This message will be returned twice by __next__.
234 event_message = self._create_event_message(ec, packet)
235
236 self._msgs = [
237 self._create_stream_beginning_message(stream),
d79a8353
SM
238 self._create_packet_beginning_message(packet),
239 event_message,
240 event_message,
241 ]
242
243 def __next__(self):
244 return self._msgs.pop(0)
245
246 class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
66964f3f 247 def __init__(self, params, obj):
d79a8353 248 tc = self._create_trace_class()
26fc5aed 249 sc = tc.create_stream_class(supports_packets=True)
d79a8353
SM
250 ec = sc.create_event_class()
251 self._add_output_port('out', (tc, sc, ec))
252
253 graph = bt2.Graph()
254 src = graph.add_component(MySource, 'src')
6c373cc9 255 it = TestOutputPortMessageIterator(graph, src.output_ports['out'])
d79a8353
SM
256
257 # Skip beginning messages.
188edac1 258 msg = next(it)
f0a42b33 259 self.assertIs(type(msg), bt2._StreamBeginningMessageConst)
188edac1 260 msg = next(it)
f0a42b33 261 self.assertIs(type(msg), bt2._PacketBeginningMessageConst)
d79a8353
SM
262
263 msg_ev1 = next(it)
264 msg_ev2 = next(it)
265
f0a42b33
FD
266 self.assertIs(type(msg_ev1), bt2._EventMessageConst)
267 self.assertIs(type(msg_ev2), bt2._EventMessageConst)
d79a8353
SM
268 self.assertEqual(msg_ev1.addr, msg_ev2.addr)
269
0a6d7302
SM
270 # Try consuming many times from an iterator that always returns TryAgain.
271 # This verifies that we are not missing an incref of Py_None, making the
272 # refcount of Py_None reach 0.
273 def test_try_again_many_times(self):
274 class MyIter(bt2._UserMessageIterator):
f00b8d40 275 def __next__(self):
0a6d7302 276 raise bt2.TryAgain
f00b8d40 277
0a6d7302 278 class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
66964f3f 279 def __init__(self, params, obj):
0a6d7302 280 self._add_output_port('out')
f00b8d40
SM
281
282 class MyFilterIter(bt2._UserMessageIterator):
283 def __init__(self, port):
284 input_port = port.user_data
ca02df0a
PP
285 self._upstream_iter = self._create_input_port_message_iterator(
286 input_port
287 )
f00b8d40
SM
288
289 def __next__(self):
290 return next(self._upstream_iter)
291
6a91742b 292 def _user_seek_beginning(self):
f00b8d40
SM
293 self._upstream_iter.seek_beginning()
294
295 @property
6a91742b 296 def _user_can_seek_beginning(self):
f00b8d40
SM
297 return self._upstream_iter.can_seek_beginning
298
299 class MyFilter(bt2._UserFilterComponent, message_iterator_class=MyFilterIter):
66964f3f 300 def __init__(self, params, obj):
f00b8d40
SM
301 input_port = self._add_input_port('in')
302 self._add_output_port('out', input_port)
303
f00b8d40
SM
304 graph = bt2.Graph()
305 src = graph.add_component(MySource, 'src')
0a6d7302
SM
306 it = TestOutputPortMessageIterator(graph, src.output_ports['out'])
307
308 # Three times the initial ref count of `None` iterations should
309 # be enough to catch the bug even if there are small differences
310 # between configurations.
311 none_ref_count = sys.getrefcount(None) * 3
312
313 for i in range(none_ref_count):
314 with self.assertRaises(bt2.TryAgain):
315 next(it)
316
317
318def _setup_seek_test(sink_cls, user_seek_beginning=None, user_can_seek_beginning=None):
319 class MySourceIter(bt2._UserMessageIterator):
320 def __init__(self, port):
321 tc, sc, ec = port.user_data
322 trace = tc()
323 stream = trace.create_stream(sc)
324 packet = stream.create_packet()
325
326 self._msgs = [
327 self._create_stream_beginning_message(stream),
328 self._create_packet_beginning_message(packet),
329 self._create_event_message(ec, packet),
330 self._create_event_message(ec, packet),
331 self._create_packet_end_message(packet),
332 self._create_stream_end_message(stream),
333 ]
334 self._at = 0
335
336 def __next__(self):
337 if self._at < len(self._msgs):
338 msg = self._msgs[self._at]
339 self._at += 1
340 return msg
341 else:
342 raise StopIteration
343
344 if user_seek_beginning is not None:
345 MySourceIter._user_seek_beginning = user_seek_beginning
346
347 if user_can_seek_beginning is not None:
348 MySourceIter._user_can_seek_beginning = property(user_can_seek_beginning)
349
350 class MySource(bt2._UserSourceComponent, message_iterator_class=MySourceIter):
351 def __init__(self, params, obj):
352 tc = self._create_trace_class()
353 sc = tc.create_stream_class(supports_packets=True)
354 ec = sc.create_event_class()
355
356 self._add_output_port('out', (tc, sc, ec))
357
358 class MyFilterIter(bt2._UserMessageIterator):
359 def __init__(self, port):
360 self._upstream_iter = self._create_input_port_message_iterator(
361 self._component._input_ports['in']
362 )
363
364 def __next__(self):
365 return next(self._upstream_iter)
366
367 @property
368 def _user_can_seek_beginning(self):
369 return self._upstream_iter.can_seek_beginning
370
371 def _user_seek_beginning(self):
372 self._upstream_iter.seek_beginning()
373
374 class MyFilter(bt2._UserFilterComponent, message_iterator_class=MyFilterIter):
375 def __init__(self, params, obj):
376 self._add_input_port('in')
377 self._add_output_port('out')
378
379 return _create_graph(MySource, sink_cls, flt_comp_cls=MyFilter)
f00b8d40 380
0a6d7302
SM
381
382class UserMessageIteratorSeekBeginningTestCase(unittest.TestCase):
f00b8d40 383 def test_can_seek_beginning(self):
6c373cc9
PP
384 class MySink(bt2._UserSinkComponent):
385 def __init__(self, params, obj):
386 self._add_input_port('in')
387
388 def _user_graph_is_configured(self):
389 self._msg_iter = self._create_input_port_message_iterator(
390 self._input_ports['in']
391 )
392
393 def _user_consume(self):
394 nonlocal can_seek_beginning
395 can_seek_beginning = self._msg_iter.can_seek_beginning
396
6a91742b 397 def _user_can_seek_beginning(self):
6c373cc9
PP
398 nonlocal input_port_iter_can_seek_beginning
399 return input_port_iter_can_seek_beginning
f00b8d40 400
0a6d7302
SM
401 graph = _setup_seek_test(
402 MySink, user_can_seek_beginning=_user_can_seek_beginning
403 )
f00b8d40 404
6c373cc9
PP
405 input_port_iter_can_seek_beginning = True
406 can_seek_beginning = None
407 graph.run_once()
408 self.assertTrue(can_seek_beginning)
f00b8d40 409
6c373cc9
PP
410 input_port_iter_can_seek_beginning = False
411 can_seek_beginning = None
412 graph.run_once()
413 self.assertFalse(can_seek_beginning)
f00b8d40 414
0a6d7302
SM
415 def test_no_can_seek_beginning_with_seek_beginning(self):
416 # Test an iterator without a _user_can_seek_beginning method, but with
417 # a _user_seek_beginning method.
418 class MySink(bt2._UserSinkComponent):
419 def __init__(self, params, obj):
420 self._add_input_port('in')
421
422 def _user_graph_is_configured(self):
423 self._msg_iter = self._create_input_port_message_iterator(
424 self._input_ports['in']
425 )
426
427 def _user_consume(self):
428 nonlocal can_seek_beginning
429 can_seek_beginning = self._msg_iter.can_seek_beginning
430
431 def _user_seek_beginning(self):
432 pass
f00b8d40 433
0a6d7302 434 graph = _setup_seek_test(MySink, user_seek_beginning=_user_seek_beginning)
6c373cc9
PP
435 can_seek_beginning = None
436 graph.run_once()
437 self.assertTrue(can_seek_beginning)
f00b8d40 438
0a6d7302
SM
439 def test_no_can_seek_beginning(self):
440 # Test an iterator without a _user_can_seek_beginning method, without
441 # a _user_seek_beginning method.
442 class MySink(bt2._UserSinkComponent):
443 def __init__(self, params, obj):
444 self._add_input_port('in')
445
446 def _user_graph_is_configured(self):
447 self._msg_iter = self._create_input_port_message_iterator(
448 self._input_ports['in']
449 )
450
451 def _user_consume(self):
452 nonlocal can_seek_beginning
453 can_seek_beginning = self._msg_iter.can_seek_beginning
454
455 graph = _setup_seek_test(MySink)
6c373cc9
PP
456 can_seek_beginning = None
457 graph.run_once()
458 self.assertFalse(can_seek_beginning)
f00b8d40
SM
459
460 def test_seek_beginning(self):
6c373cc9
PP
461 class MySink(bt2._UserSinkComponent):
462 def __init__(self, params, obj):
463 self._add_input_port('in')
f00b8d40 464
6c373cc9
PP
465 def _user_graph_is_configured(self):
466 self._msg_iter = self._create_input_port_message_iterator(
467 self._input_ports['in']
468 )
469
470 def _user_consume(self):
471 nonlocal do_seek_beginning
472 nonlocal msg
473
474 if do_seek_beginning:
475 self._msg_iter.seek_beginning()
476 return
477
478 msg = next(self._msg_iter)
479
0a6d7302
SM
480 def _user_seek_beginning(self):
481 self._at = 0
482
6c373cc9 483 msg = None
0a6d7302
SM
484 graph = _setup_seek_test(MySink, user_seek_beginning=_user_seek_beginning)
485
486 # Consume message.
487 do_seek_beginning = False
6c373cc9 488 graph.run_once()
f0a42b33 489 self.assertIs(type(msg), bt2._StreamBeginningMessageConst)
0a6d7302
SM
490
491 # Consume message.
6c373cc9 492 graph.run_once()
f0a42b33 493 self.assertIs(type(msg), bt2._PacketBeginningMessageConst)
0a6d7302
SM
494
495 # Seek beginning.
6c373cc9
PP
496 do_seek_beginning = True
497 graph.run_once()
0a6d7302
SM
498
499 # Consume message.
6c373cc9
PP
500 do_seek_beginning = False
501 graph.run_once()
f0a42b33 502 self.assertIs(type(msg), bt2._StreamBeginningMessageConst)
f00b8d40 503
6c373cc9
PP
504 def test_seek_beginning_user_error(self):
505 class MySink(bt2._UserSinkComponent):
506 def __init__(self, params, obj):
507 self._add_input_port('in')
f00b8d40 508
6c373cc9
PP
509 def _user_graph_is_configured(self):
510 self._msg_iter = self._create_input_port_message_iterator(
511 self._input_ports['in']
512 )
f00b8d40 513
6c373cc9
PP
514 def _user_consume(self):
515 self._msg_iter.seek_beginning()
f00b8d40 516
0a6d7302 517 def _user_seek_beginning(self):
cfbd7cf3 518 raise ValueError('ouch')
f00b8d40 519
0a6d7302 520 graph = _setup_seek_test(MySink, user_seek_beginning=_user_seek_beginning)
f00b8d40 521
694c792b 522 with self.assertRaises(bt2._Error):
6c373cc9 523 graph.run_once()
f00b8d40
SM
524
525
f00b8d40
SM
526if __name__ == '__main__':
527 unittest.main()
This page took 0.06779 seconds and 4 git commands to generate.