tests/bindings/python/bt2/test_message_iterator.py

   1 # SPDX-License-Identifier: GPL-2.0-only
   2 #
   3 # Copyright (C) 2019 EfficiOS Inc.
   4 #
   5
   6 import unittest
   7 import bt2
   8 import sys
   9 from utils import TestOutputPortMessageIterator
  10 from bt2 import port as bt2_port
  11 from bt2 import message_iterator as bt2_message_iterator
  12
  13
  14 class SimpleSink(bt2._UserSinkComponent):
  15     # Straightforward sink that creates one input port (`in`) and consumes from
  16     # it.
  17
  18     def __init__(self, config, params, obj):
  19         self._add_input_port('in')
  20
  21     def _user_consume(self):
  22         next(self._msg_iter)
  23
  24     def _user_graph_is_configured(self):
  25         self._msg_iter = self._create_message_iterator(self._input_ports['in'])
  26
  27
  28 def _create_graph(src_comp_cls, sink_comp_cls, flt_comp_cls=None):
  29     graph = bt2.Graph()
  30
  31     src_comp = graph.add_component(src_comp_cls, 'src')
  32     sink_comp = graph.add_component(sink_comp_cls, 'sink')
  33
  34     if flt_comp_cls is not None:
  35         flt_comp = graph.add_component(flt_comp_cls, 'flt')
  36         graph.connect_ports(src_comp.output_ports['out'], flt_comp.input_ports['in'])
  37         graph.connect_ports(flt_comp.output_ports['out'], sink_comp.input_ports['in'])
  38     else:
  39         graph.connect_ports(src_comp.output_ports['out'], sink_comp.input_ports['in'])
  40
  41     return graph
  42
  43
  44 class UserMessageIteratorTestCase(unittest.TestCase):
  45     def test_init(self):
  46         the_output_port_from_source = None
  47         the_output_port_from_iter = None
  48
  49         class MyIter(bt2._UserMessageIterator):
  50             def __init__(self, config, self_port_output):
  51                 nonlocal initialized
  52                 nonlocal the_output_port_from_iter
  53                 initialized = True
  54                 the_output_port_from_iter = self_port_output
  55
  56         class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
  57             def __init__(self, config, params, obj):
  58                 nonlocal the_output_port_from_source
  59                 the_output_port_from_source = self._add_output_port('out', 'user data')
  60
  61         initialized = False
  62         graph = _create_graph(MySource, SimpleSink)
  63         graph.run()
  64         self.assertTrue(initialized)
  65         self.assertEqual(
  66             the_output_port_from_source.addr, the_output_port_from_iter.addr
  67         )
  68         self.assertEqual(the_output_port_from_iter.user_data, 'user data')
  69
  70     def test_create_from_message_iterator(self):
  71         class MySourceIter(bt2._UserMessageIterator):
  72             def __init__(self, config, self_port_output):
  73                 nonlocal src_iter_initialized
  74                 src_iter_initialized = True
  75
  76         class MySource(bt2._UserSourceComponent, message_iterator_class=MySourceIter):
  77             def __init__(self, config, params, obj):
  78                 self._add_output_port('out')
  79
  80         class MyFilterIter(bt2._UserMessageIterator):
  81             def __init__(self, config, self_port_output):
  82                 nonlocal flt_iter_initialized
  83                 flt_iter_initialized = True
  84                 self._up_iter = self._create_message_iterator(
  85                     self._component._input_ports['in']
  86                 )
  87
  88             def __next__(self):
  89                 return next(self._up_iter)
  90
  91         class MyFilter(bt2._UserFilterComponent, message_iterator_class=MyFilterIter):
  92             def __init__(self, config, params, obj):
  93                 self._add_input_port('in')
  94                 self._add_output_port('out')
  95
  96         src_iter_initialized = False
  97         flt_iter_initialized = False
  98         graph = _create_graph(MySource, SimpleSink, MyFilter)
  99         graph.run()
 100         self.assertTrue(src_iter_initialized)
 101         self.assertTrue(flt_iter_initialized)
 102
 103     def test_create_user_error(self):
 104         # This tests both error handling by
 105         # _UserSinkComponent._create_message_iterator
 106         # and _UserMessageIterator._create_message_iterator, as they
 107         # are both used in the graph.
 108         class MySourceIter(bt2._UserMessageIterator):
 109             def __init__(self, config, self_port_output):
 110                 raise ValueError('Very bad error')
 111
 112         class MySource(bt2._UserSourceComponent, message_iterator_class=MySourceIter):
 113             def __init__(self, config, params, obj):
 114                 self._add_output_port('out')
 115
 116         class MyFilterIter(bt2._UserMessageIterator):
 117             def __init__(self, config, self_port_output):
 118                 # This is expected to raise because of the error in
 119                 # MySourceIter.__init__.
 120                 self._create_message_iterator(self._component._input_ports['in'])
 121
 122         class MyFilter(bt2._UserFilterComponent, message_iterator_class=MyFilterIter):
 123             def __init__(self, config, params, obj):
 124                 self._add_input_port('in')
 125                 self._add_output_port('out')
 126
 127         graph = _create_graph(MySource, SimpleSink, MyFilter)
 128
 129         with self.assertRaises(bt2._Error) as ctx:
 130             graph.run()
 131
 132         exc = ctx.exception
 133         cause = exc[0]
 134
 135         self.assertIsInstance(cause, bt2._MessageIteratorErrorCause)
 136         self.assertEqual(cause.component_name, 'src')
 137         self.assertEqual(cause.component_output_port_name, 'out')
 138         self.assertIn('ValueError: Very bad error', cause.message)
 139
 140     def test_finalize(self):
 141         class MyIter(bt2._UserMessageIterator):
 142             def _user_finalize(self):
 143                 nonlocal finalized
 144                 finalized = True
 145
 146         class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
 147             def __init__(self, config, params, obj):
 148                 self._add_output_port('out')
 149
 150         finalized = False
 151         graph = _create_graph(MySource, SimpleSink)
 152         graph.run()
 153         del graph
 154         self.assertTrue(finalized)
 155
 156     def test_config_parameter(self):
 157         class MyIter(bt2._UserMessageIterator):
 158             def __init__(self, config, port):
 159                 nonlocal config_type
 160                 config_type = type(config)
 161
 162         class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
 163             def __init__(self, config, params, obj):
 164                 self._add_output_port('out')
 165
 166         config_type = None
 167         graph = _create_graph(MySource, SimpleSink)
 168         graph.run()
 169         self.assertIs(config_type, bt2_message_iterator._MessageIteratorConfiguration)
 170
 171     def _test_config_can_seek_forward(self, set_can_seek_forward):
 172         class MyIter(bt2._UserMessageIterator):
 173             def __init__(self, config, port):
 174                 if set_can_seek_forward:
 175                     config.can_seek_forward = True
 176
 177         class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
 178             def __init__(self, config, params, obj):
 179                 self._add_output_port('out')
 180
 181         class MySink(bt2._UserSinkComponent):
 182             def __init__(self, config, params, obj):
 183                 self._add_input_port('in')
 184
 185             def _user_graph_is_configured(self):
 186                 self._msg_iter = self._create_message_iterator(self._input_ports['in'])
 187
 188             def _user_consume(self):
 189                 nonlocal can_seek_forward
 190                 can_seek_forward = self._msg_iter.can_seek_forward
 191
 192         can_seek_forward = None
 193         graph = _create_graph(MySource, MySink)
 194         graph.run_once()
 195         self.assertIs(can_seek_forward, set_can_seek_forward)
 196
 197     def test_config_can_seek_forward_default(self):
 198         self._test_config_can_seek_forward(False)
 199
 200     def test_config_can_seek_forward(self):
 201         self._test_config_can_seek_forward(True)
 202
 203     def test_config_can_seek_forward_wrong_type(self):
 204         class MyIter(bt2._UserMessageIterator):
 205             def __init__(self, config, port):
 206                 config.can_seek_forward = 1
 207
 208         class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
 209             def __init__(self, config, params, obj):
 210                 self._add_output_port('out')
 211
 212         graph = _create_graph(MySource, SimpleSink)
 213         with self.assertRaises(bt2._Error) as ctx:
 214             graph.run()
 215
 216         root_cause = ctx.exception[0]
 217         self.assertIn("TypeError: 'int' is not a 'bool' object", root_cause.message)
 218
 219     def test_component(self):
 220         class MyIter(bt2._UserMessageIterator):
 221             def __init__(self, config, self_port_output):
 222                 nonlocal salut
 223                 salut = self._component._salut
 224
 225         class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
 226             def __init__(self, config, params, obj):
 227                 self._add_output_port('out')
 228                 self._salut = 23
 229
 230         salut = None
 231         graph = _create_graph(MySource, SimpleSink)
 232         graph.run()
 233         self.assertEqual(salut, 23)
 234
 235     def test_port(self):
 236         class MyIter(bt2._UserMessageIterator):
 237             def __init__(self_iter, config, self_port_output):
 238                 nonlocal called
 239                 called = True
 240                 port = self_iter._port
 241                 self.assertIs(type(self_port_output), bt2_port._UserComponentOutputPort)
 242                 self.assertIs(type(port), bt2_port._UserComponentOutputPort)
 243                 self.assertEqual(self_port_output.addr, port.addr)
 244
 245         class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
 246             def __init__(self, config, params, obj):
 247                 self._add_output_port('out')
 248
 249         called = False
 250         graph = _create_graph(MySource, SimpleSink)
 251         graph.run()
 252         self.assertTrue(called)
 253
 254     def test_addr(self):
 255         class MyIter(bt2._UserMessageIterator):
 256             def __init__(self, config, self_port_output):
 257                 nonlocal addr
 258                 addr = self.addr
 259
 260         class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
 261             def __init__(self, config, params, obj):
 262                 self._add_output_port('out')
 263
 264         addr = None
 265         graph = _create_graph(MySource, SimpleSink)
 266         graph.run()
 267         self.assertIsNotNone(addr)
 268         self.assertNotEqual(addr, 0)
 269
 270     # Test that messages returned by _UserMessageIterator.__next__ remain valid
 271     # and can be re-used.
 272     def test_reuse_message(self):
 273         class MyIter(bt2._UserMessageIterator):
 274             def __init__(self, config, port):
 275                 tc, sc, ec = port.user_data
 276                 trace = tc()
 277                 stream = trace.create_stream(sc)
 278                 packet = stream.create_packet()
 279
 280                 # This message will be returned twice by __next__.
 281                 event_message = self._create_event_message(ec, packet)
 282
 283                 self._msgs = [
 284                     self._create_stream_beginning_message(stream),
 285                     self._create_packet_beginning_message(packet),
 286                     event_message,
 287                     event_message,
 288                 ]
 289
 290             def __next__(self):
 291                 return self._msgs.pop(0)
 292
 293         class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
 294             def __init__(self, config, params, obj):
 295                 tc = self._create_trace_class()
 296                 sc = tc.create_stream_class(supports_packets=True)
 297                 ec = sc.create_event_class()
 298                 self._add_output_port('out', (tc, sc, ec))
 299
 300         graph = bt2.Graph()
 301         src = graph.add_component(MySource, 'src')
 302         it = TestOutputPortMessageIterator(graph, src.output_ports['out'])
 303
 304         # Skip beginning messages.
 305         msg = next(it)
 306         self.assertIs(type(msg), bt2._StreamBeginningMessageConst)
 307         msg = next(it)
 308         self.assertIs(type(msg), bt2._PacketBeginningMessageConst)
 309
 310         msg_ev1 = next(it)
 311         msg_ev2 = next(it)
 312
 313         self.assertIs(type(msg_ev1), bt2._EventMessageConst)
 314         self.assertIs(type(msg_ev2), bt2._EventMessageConst)
 315         self.assertEqual(msg_ev1.addr, msg_ev2.addr)
 316
 317     # Try consuming many times from an iterator that always returns TryAgain.
 318     # This verifies that we are not missing an incref of Py_None, making the
 319     # refcount of Py_None reach 0.
 320     def test_try_again_many_times(self):
 321         class MyIter(bt2._UserMessageIterator):
 322             def __next__(self):
 323                 raise bt2.TryAgain
 324
 325         class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
 326             def __init__(self, config, params, obj):
 327                 self._add_output_port('out')
 328
 329         class MyFilterIter(bt2._UserMessageIterator):
 330             def __init__(self, port):
 331                 input_port = port.user_data
 332                 self._upstream_iter = self._create_message_iterator(input_port)
 333
 334             def __next__(self):
 335                 return next(self._upstream_iter)
 336
 337             def _user_seek_beginning(self):
 338                 self._upstream_iter.seek_beginning()
 339
 340             def _user_can_seek_beginning(self):
 341                 return self._upstream_iter.can_seek_beginning()
 342
 343         class MyFilter(bt2._UserFilterComponent, message_iterator_class=MyFilterIter):
 344             def __init__(self, config, params, obj):
 345                 input_port = self._add_input_port('in')
 346                 self._add_output_port('out', input_port)
 347
 348         graph = bt2.Graph()
 349         src = graph.add_component(MySource, 'src')
 350         it = TestOutputPortMessageIterator(graph, src.output_ports['out'])
 351
 352         # Three times the initial ref count of `None` iterations should
 353         # be enough to catch the bug even if there are small differences
 354         # between configurations.
 355         none_ref_count = sys.getrefcount(None) * 3
 356
 357         for i in range(none_ref_count):
 358             with self.assertRaises(bt2.TryAgain):
 359                 next(it)
 360
 361     def test_error_in_iterator_with_cycle_after_having_created_upstream_iterator(self):
 362         # Test a failure that triggered an abort in libbabeltrace2, in this situation:
 363         #
 364         #   - The filter iterator creates an upstream iterator.
 365         #   - The filter iterator creates a reference cycle, including itself.
 366         #   - An exception is raised, causing the filter iterator's
 367         #     initialization method to fail.
 368         class MySourceIter(bt2._UserMessageIterator):
 369             pass
 370
 371         class MySource(bt2._UserSourceComponent, message_iterator_class=MySourceIter):
 372             def __init__(self, config, params, obj):
 373                 self._add_output_port('out')
 374
 375         class MyFilterIter(bt2._UserMessageIterator):
 376             def __init__(self, config, port):
 377                 # First, create an upstream iterator.
 378                 self._upstream_iter = self._create_message_iterator(
 379                     self._component._input_ports['in']
 380                 )
 381
 382                 # Then, voluntarily make a reference cycle that will keep this
 383                 # Python object alive, which will keep the upstream iterator
 384                 # Babeltrace object alive.
 385                 self._self = self
 386
 387                 # Finally, raise an exception to make __init__ fail.
 388                 raise ValueError('woops')
 389
 390         class MyFilter(bt2._UserFilterComponent, message_iterator_class=MyFilterIter):
 391             def __init__(self, config, params, obj):
 392                 self._in = self._add_input_port('in')
 393                 self._out = self._add_output_port('out')
 394
 395         class MySink(bt2._UserSinkComponent):
 396             def __init__(self, config, params, obj):
 397                 self._input_port = self._add_input_port('in')
 398
 399             def _user_graph_is_configured(self):
 400                 self._upstream_iter = self._create_message_iterator(self._input_port)
 401
 402             def _user_consume(self):
 403                 # We should not reach this.
 404                 assert False
 405
 406         g = bt2.Graph()
 407         src = g.add_component(MySource, 'src')
 408         flt = g.add_component(MyFilter, 'flt')
 409         snk = g.add_component(MySink, 'snk')
 410         g.connect_ports(src.output_ports['out'], flt.input_ports['in'])
 411         g.connect_ports(flt.output_ports['out'], snk.input_ports['in'])
 412
 413         with self.assertRaisesRegex(bt2._Error, 'ValueError: woops'):
 414             g.run()
 415
 416
 417 def _setup_seek_test(
 418     sink_cls,
 419     user_seek_beginning=None,
 420     user_can_seek_beginning=None,
 421     user_seek_ns_from_origin=None,
 422     user_can_seek_ns_from_origin=None,
 423     can_seek_forward=False,
 424 ):
 425     class MySourceIter(bt2._UserMessageIterator):
 426         def __init__(self, config, port):
 427             tc, sc, ec = port.user_data
 428             trace = tc()
 429             stream = trace.create_stream(sc)
 430             packet = stream.create_packet()
 431
 432             self._msgs = [
 433                 self._create_stream_beginning_message(stream),
 434                 self._create_packet_beginning_message(packet),
 435                 self._create_event_message(ec, packet),
 436                 self._create_event_message(ec, packet),
 437                 self._create_packet_end_message(packet),
 438                 self._create_stream_end_message(stream),
 439             ]
 440             self._at = 0
 441             config.can_seek_forward = can_seek_forward
 442
 443         def __next__(self):
 444             if self._at < len(self._msgs):
 445                 msg = self._msgs[self._at]
 446                 self._at += 1
 447                 return msg
 448             else:
 449                 raise StopIteration
 450
 451     if user_seek_beginning is not None:
 452         MySourceIter._user_seek_beginning = user_seek_beginning
 453
 454     if user_can_seek_beginning is not None:
 455         MySourceIter._user_can_seek_beginning = user_can_seek_beginning
 456
 457     if user_seek_ns_from_origin is not None:
 458         MySourceIter._user_seek_ns_from_origin = user_seek_ns_from_origin
 459
 460     if user_can_seek_ns_from_origin is not None:
 461         MySourceIter._user_can_seek_ns_from_origin = user_can_seek_ns_from_origin
 462
 463     class MySource(bt2._UserSourceComponent, message_iterator_class=MySourceIter):
 464         def __init__(self, config, params, obj):
 465             tc = self._create_trace_class()
 466             sc = tc.create_stream_class(supports_packets=True)
 467             ec = sc.create_event_class()
 468
 469             self._add_output_port('out', (tc, sc, ec))
 470
 471     class MyFilterIter(bt2._UserMessageIterator):
 472         def __init__(self, config, port):
 473             self._upstream_iter = self._create_message_iterator(
 474                 self._component._input_ports['in']
 475             )
 476             config.can_seek_forward = self._upstream_iter.can_seek_forward
 477
 478         def __next__(self):
 479             return next(self._upstream_iter)
 480
 481         def _user_can_seek_beginning(self):
 482             return self._upstream_iter.can_seek_beginning()
 483
 484         def _user_seek_beginning(self):
 485             self._upstream_iter.seek_beginning()
 486
 487         def _user_can_seek_ns_from_origin(self, ns_from_origin):
 488             return self._upstream_iter.can_seek_ns_from_origin(ns_from_origin)
 489
 490         def _user_seek_ns_from_origin(self, ns_from_origin):
 491             self._upstream_iter.seek_ns_from_origin(ns_from_origin)
 492
 493     class MyFilter(bt2._UserFilterComponent, message_iterator_class=MyFilterIter):
 494         def __init__(self, config, params, obj):
 495             self._add_input_port('in')
 496             self._add_output_port('out')
 497
 498     return _create_graph(MySource, sink_cls, flt_comp_cls=MyFilter)
 499
 500
 501 class UserMessageIteratorSeekBeginningTestCase(unittest.TestCase):
 502     def test_can_seek_beginning_without_seek_beginning(self):
 503         with self.assertRaisesRegex(
 504             bt2._IncompleteUserClass,
 505             "cannot create component class 'MySource': message iterator class implements _user_can_seek_beginning but not _user_seek_beginning",
 506         ):
 507             _setup_seek_test(SimpleSink, user_can_seek_beginning=lambda: None)
 508
 509     def test_can_seek_beginning(self):
 510         class MySink(bt2._UserSinkComponent):
 511             def __init__(self, config, params, obj):
 512                 self._add_input_port('in')
 513
 514             def _user_graph_is_configured(self):
 515                 self._msg_iter = self._create_message_iterator(self._input_ports['in'])
 516
 517             def _user_consume(self):
 518                 nonlocal can_seek_beginning
 519                 can_seek_beginning = self._msg_iter.can_seek_beginning()
 520
 521         def _user_can_seek_beginning(self):
 522             nonlocal input_port_iter_can_seek_beginning
 523             return input_port_iter_can_seek_beginning
 524
 525         graph = _setup_seek_test(
 526             MySink,
 527             user_can_seek_beginning=_user_can_seek_beginning,
 528             user_seek_beginning=lambda: None,
 529         )
 530
 531         input_port_iter_can_seek_beginning = True
 532         can_seek_beginning = None
 533         graph.run_once()
 534         self.assertIs(can_seek_beginning, True)
 535
 536         input_port_iter_can_seek_beginning = False
 537         can_seek_beginning = None
 538         graph.run_once()
 539         self.assertIs(can_seek_beginning, False)
 540
 541     def test_no_can_seek_beginning_with_seek_beginning(self):
 542         # Test an iterator without a _user_can_seek_beginning method, but with
 543         # a _user_seek_beginning method.
 544         class MySink(bt2._UserSinkComponent):
 545             def __init__(self, config, params, obj):
 546                 self._add_input_port('in')
 547
 548             def _user_graph_is_configured(self):
 549                 self._msg_iter = self._create_message_iterator(self._input_ports['in'])
 550
 551             def _user_consume(self):
 552                 nonlocal can_seek_beginning
 553                 can_seek_beginning = self._msg_iter.can_seek_beginning()
 554
 555         def _user_seek_beginning(self):
 556             pass
 557
 558         graph = _setup_seek_test(MySink, user_seek_beginning=_user_seek_beginning)
 559         can_seek_beginning = None
 560         graph.run_once()
 561         self.assertIs(can_seek_beginning, True)
 562
 563     def test_no_can_seek_beginning(self):
 564         # Test an iterator without a _user_can_seek_beginning method, without
 565         # a _user_seek_beginning method.
 566         class MySink(bt2._UserSinkComponent):
 567             def __init__(self, config, params, obj):
 568                 self._add_input_port('in')
 569
 570             def _user_graph_is_configured(self):
 571                 self._msg_iter = self._create_message_iterator(self._input_ports['in'])
 572
 573             def _user_consume(self):
 574                 nonlocal can_seek_beginning
 575                 can_seek_beginning = self._msg_iter.can_seek_beginning()
 576
 577         graph = _setup_seek_test(MySink)
 578         can_seek_beginning = None
 579         graph.run_once()
 580         self.assertIs(can_seek_beginning, False)
 581
 582     def test_can_seek_beginning_user_error(self):
 583         class MySink(bt2._UserSinkComponent):
 584             def __init__(self, config, params, obj):
 585                 self._add_input_port('in')
 586
 587             def _user_graph_is_configured(self):
 588                 self._msg_iter = self._create_message_iterator(self._input_ports['in'])
 589
 590             def _user_consume(self):
 591                 # This is expected to raise.
 592                 self._msg_iter.can_seek_beginning()
 593
 594         def _user_can_seek_beginning(self):
 595             raise ValueError('moustiquaire')
 596
 597         graph = _setup_seek_test(
 598             MySink,
 599             user_can_seek_beginning=_user_can_seek_beginning,
 600             user_seek_beginning=lambda: None,
 601         )
 602
 603         with self.assertRaises(bt2._Error) as ctx:
 604             graph.run_once()
 605
 606         cause = ctx.exception[0]
 607         self.assertIn('ValueError: moustiquaire', cause.message)
 608
 609     def test_can_seek_beginning_wrong_return_value(self):
 610         class MySink(bt2._UserSinkComponent):
 611             def __init__(self, config, params, obj):
 612                 self._add_input_port('in')
 613
 614             def _user_graph_is_configured(self):
 615                 self._msg_iter = self._create_message_iterator(self._input_ports['in'])
 616
 617             def _user_consume(self):
 618                 # This is expected to raise.
 619                 self._msg_iter.can_seek_beginning()
 620
 621         def _user_can_seek_beginning(self):
 622             return 'Amqui'
 623
 624         graph = _setup_seek_test(
 625             MySink,
 626             user_can_seek_beginning=_user_can_seek_beginning,
 627             user_seek_beginning=lambda: None,
 628         )
 629
 630         with self.assertRaises(bt2._Error) as ctx:
 631             graph.run_once()
 632
 633         cause = ctx.exception[0]
 634         self.assertIn("TypeError: 'str' is not a 'bool' object", cause.message)
 635
 636     def test_seek_beginning(self):
 637         class MySink(bt2._UserSinkComponent):
 638             def __init__(self, config, params, obj):
 639                 self._add_input_port('in')
 640
 641             def _user_graph_is_configured(self):
 642                 self._msg_iter = self._create_message_iterator(self._input_ports['in'])
 643
 644             def _user_consume(self):
 645                 nonlocal do_seek_beginning
 646                 nonlocal msg
 647
 648                 if do_seek_beginning:
 649                     self._msg_iter.seek_beginning()
 650                     return
 651
 652                 msg = next(self._msg_iter)
 653
 654         def _user_seek_beginning(self):
 655             self._at = 0
 656
 657         msg = None
 658         graph = _setup_seek_test(MySink, user_seek_beginning=_user_seek_beginning)
 659
 660         # Consume message.
 661         do_seek_beginning = False
 662         graph.run_once()
 663         self.assertIs(type(msg), bt2._StreamBeginningMessageConst)
 664
 665         # Consume message.
 666         graph.run_once()
 667         self.assertIs(type(msg), bt2._PacketBeginningMessageConst)
 668
 669         # Seek beginning.
 670         do_seek_beginning = True
 671         graph.run_once()
 672
 673         # Consume message.
 674         do_seek_beginning = False
 675         graph.run_once()
 676         self.assertIs(type(msg), bt2._StreamBeginningMessageConst)
 677
 678     def test_seek_beginning_user_error(self):
 679         class MySink(bt2._UserSinkComponent):
 680             def __init__(self, config, params, obj):
 681                 self._add_input_port('in')
 682
 683             def _user_graph_is_configured(self):
 684                 self._msg_iter = self._create_message_iterator(self._input_ports['in'])
 685
 686             def _user_consume(self):
 687                 self._msg_iter.seek_beginning()
 688
 689         def _user_seek_beginning(self):
 690             raise ValueError('ouch')
 691
 692         graph = _setup_seek_test(MySink, user_seek_beginning=_user_seek_beginning)
 693
 694         with self.assertRaises(bt2._Error):
 695             graph.run_once()
 696
 697
 698 class UserMessageIteratorSeekNsFromOriginTestCase(unittest.TestCase):
 699     def test_can_seek_ns_from_origin_without_seek_ns_from_origin(self):
 700         # Test the case where:
 701         #
 702         #   - can_seek_ns_from_origin: Returns True (don't really care, as long
 703         #     as it's provided)
 704         #   - seek_ns_from_origin provided: No
 705         #   - can the iterator seek beginning: Don't care
 706         #   - can the iterator seek forward: Don't care
 707         for can_seek_ns_from_origin in (False, True):
 708             for iter_can_seek_beginning in (False, True):
 709                 for iter_can_seek_forward in (False, True):
 710                     with self.assertRaisesRegex(
 711                         bt2._IncompleteUserClass,
 712                         "cannot create component class 'MySource': message iterator class implements _user_can_seek_ns_from_origin but not _user_seek_ns_from_origin",
 713                     ):
 714                         self._can_seek_ns_from_origin_test(
 715                             None,
 716                             user_can_seek_ns_from_origin_ret_val=True,
 717                             user_seek_ns_from_origin_provided=False,
 718                             iter_can_seek_beginning=iter_can_seek_beginning,
 719                             iter_can_seek_forward=iter_can_seek_forward,
 720                         )
 721
 722     def test_can_seek_ns_from_origin_returns_true(self):
 723         # Test the case where:
 724         #
 725         #   - can_seek_ns_from_origin: returns True
 726         #   - seek_ns_from_origin provided: Yes
 727         #   - can the iterator seek beginning: Don't care
 728         #   - can the iterator seek forward: Don't care
 729         #
 730         # We expect iter.can_seek_ns_from_origin to return True.
 731         for iter_can_seek_beginning in (False, True):
 732             for iter_can_seek_forward in (False, True):
 733                 self._can_seek_ns_from_origin_test(
 734                     expected_outcome=True,
 735                     user_can_seek_ns_from_origin_ret_val=True,
 736                     user_seek_ns_from_origin_provided=True,
 737                     iter_can_seek_beginning=iter_can_seek_beginning,
 738                     iter_can_seek_forward=iter_can_seek_forward,
 739                 )
 740
 741     def test_can_seek_ns_from_origin_returns_false_can_seek_beginning_forward_seekable(
 742         self,
 743     ):
 744         # Test the case where:
 745         #
 746         #   - can_seek_ns_from_origin: returns False
 747         #   - seek_ns_from_origin provided: Yes
 748         #   - can the iterator seek beginning: Yes
 749         #   - can the iterator seek forward: Yes
 750         #
 751         # We expect iter.can_seek_ns_from_origin to return True.
 752         self._can_seek_ns_from_origin_test(
 753             expected_outcome=True,
 754             user_can_seek_ns_from_origin_ret_val=False,
 755             user_seek_ns_from_origin_provided=True,
 756             iter_can_seek_beginning=True,
 757             iter_can_seek_forward=True,
 758         )
 759
 760     def test_can_seek_ns_from_origin_returns_false_can_seek_beginning_not_forward_seekable(
 761         self,
 762     ):
 763         # Test the case where:
 764         #
 765         #   - can_seek_ns_from_origin: returns False
 766         #   - seek_ns_from_origin provided: Yes
 767         #   - can the iterator seek beginning: Yes
 768         #   - can the iterator seek forward: No
 769         #
 770         # We expect iter.can_seek_ns_from_origin to return False.
 771         self._can_seek_ns_from_origin_test(
 772             expected_outcome=False,
 773             user_can_seek_ns_from_origin_ret_val=False,
 774             user_seek_ns_from_origin_provided=True,
 775             iter_can_seek_beginning=True,
 776             iter_can_seek_forward=False,
 777         )
 778
 779     def test_can_seek_ns_from_origin_returns_false_cant_seek_beginning_forward_seekable(
 780         self,
 781     ):
 782         # Test the case where:
 783         #
 784         #   - can_seek_ns_from_origin: returns False
 785         #   - seek_ns_from_origin provided: Yes
 786         #   - can the iterator seek beginning: No
 787         #   - can the iterator seek forward: Yes
 788         #
 789         # We expect iter.can_seek_ns_from_origin to return False.
 790         self._can_seek_ns_from_origin_test(
 791             expected_outcome=False,
 792             user_can_seek_ns_from_origin_ret_val=False,
 793             user_seek_ns_from_origin_provided=True,
 794             iter_can_seek_beginning=False,
 795             iter_can_seek_forward=True,
 796         )
 797
 798     def test_can_seek_ns_from_origin_returns_false_cant_seek_beginning_not_forward_seekable(
 799         self,
 800     ):
 801         # Test the case where:
 802         #
 803         #   - can_seek_ns_from_origin: returns False
 804         #   - seek_ns_from_origin provided: Yes
 805         #   - can the iterator seek beginning: No
 806         #   - can the iterator seek forward: No
 807         #
 808         # We expect iter.can_seek_ns_from_origin to return False.
 809         self._can_seek_ns_from_origin_test(
 810             expected_outcome=False,
 811             user_can_seek_ns_from_origin_ret_val=False,
 812             user_seek_ns_from_origin_provided=True,
 813             iter_can_seek_beginning=False,
 814             iter_can_seek_forward=False,
 815         )
 816
 817     def test_no_can_seek_ns_from_origin_seek_ns_from_origin(self):
 818         # Test the case where:
 819         #
 820         #   - can_seek_ns_from_origin: Not provided
 821         #   - seek_ns_from_origin provided: Yes
 822         #   - can the iterator seek beginning: Don't care
 823         #   - can the iterator seek forward: Don't care
 824         #
 825         # We expect iter.can_seek_ns_from_origin to return True.
 826         for iter_can_seek_beginning in (False, True):
 827             for iter_can_seek_forward in (False, True):
 828                 self._can_seek_ns_from_origin_test(
 829                     expected_outcome=True,
 830                     user_can_seek_ns_from_origin_ret_val=None,
 831                     user_seek_ns_from_origin_provided=True,
 832                     iter_can_seek_beginning=iter_can_seek_beginning,
 833                     iter_can_seek_forward=iter_can_seek_forward,
 834                 )
 835
 836     def test_no_can_seek_ns_from_origin_no_seek_ns_from_origin_can_seek_beginning_forward_seekable(
 837         self,
 838     ):
 839         # Test the case where:
 840         #
 841         #   - can_seek_ns_from_origin: Not provided
 842         #   - seek_ns_from_origin provided: Not provided
 843         #   - can the iterator seek beginning: Yes
 844         #   - can the iterator seek forward: Yes
 845         #
 846         # We expect iter.can_seek_ns_from_origin to return True.
 847         self._can_seek_ns_from_origin_test(
 848             expected_outcome=True,
 849             user_can_seek_ns_from_origin_ret_val=None,
 850             user_seek_ns_from_origin_provided=False,
 851             iter_can_seek_beginning=True,
 852             iter_can_seek_forward=True,
 853         )
 854
 855     def test_no_can_seek_ns_from_origin_no_seek_ns_from_origin_can_seek_beginning_not_forward_seekable(
 856         self,
 857     ):
 858         # Test the case where:
 859         #
 860         #   - can_seek_ns_from_origin: Not provided
 861         #   - seek_ns_from_origin provided: Not provided
 862         #   - can the iterator seek beginning: Yes
 863         #   - can the iterator seek forward: No
 864         #
 865         # We expect iter.can_seek_ns_from_origin to return False.
 866         self._can_seek_ns_from_origin_test(
 867             expected_outcome=False,
 868             user_can_seek_ns_from_origin_ret_val=None,
 869             user_seek_ns_from_origin_provided=False,
 870             iter_can_seek_beginning=True,
 871             iter_can_seek_forward=False,
 872         )
 873
 874     def test_no_can_seek_ns_from_origin_no_seek_ns_from_origin_cant_seek_beginning_forward_seekable(
 875         self,
 876     ):
 877         # Test the case where:
 878         #
 879         #   - can_seek_ns_from_origin: Not provided
 880         #   - seek_ns_from_origin provided: Not provided
 881         #   - can the iterator seek beginning: No
 882         #   - can the iterator seek forward: Yes
 883         #
 884         # We expect iter.can_seek_ns_from_origin to return False.
 885         self._can_seek_ns_from_origin_test(
 886             expected_outcome=False,
 887             user_can_seek_ns_from_origin_ret_val=None,
 888             user_seek_ns_from_origin_provided=False,
 889             iter_can_seek_beginning=False,
 890             iter_can_seek_forward=True,
 891         )
 892
 893     def test_no_can_seek_ns_from_origin_no_seek_ns_from_origin_cant_seek_beginning_not_forward_seekable(
 894         self,
 895     ):
 896         # Test the case where:
 897         #
 898         #   - can_seek_ns_from_origin: Not provided
 899         #   - seek_ns_from_origin provided: Not provided
 900         #   - can the iterator seek beginning: No
 901         #   - can the iterator seek forward: No
 902         #
 903         # We expect iter.can_seek_ns_from_origin to return False.
 904         self._can_seek_ns_from_origin_test(
 905             expected_outcome=False,
 906             user_can_seek_ns_from_origin_ret_val=None,
 907             user_seek_ns_from_origin_provided=False,
 908             iter_can_seek_beginning=False,
 909             iter_can_seek_forward=False,
 910         )
 911
 912     def _can_seek_ns_from_origin_test(
 913         self,
 914         expected_outcome,
 915         user_can_seek_ns_from_origin_ret_val,
 916         user_seek_ns_from_origin_provided,
 917         iter_can_seek_beginning,
 918         iter_can_seek_forward,
 919     ):
 920         class MySink(bt2._UserSinkComponent):
 921             def __init__(self, config, params, obj):
 922                 self._add_input_port('in')
 923
 924             def _user_graph_is_configured(self):
 925                 self._msg_iter = self._create_message_iterator(self._input_ports['in'])
 926
 927             def _user_consume(self):
 928                 nonlocal can_seek_ns_from_origin
 929                 can_seek_ns_from_origin = self._msg_iter.can_seek_ns_from_origin(
 930                     passed_ns_from_origin
 931                 )
 932
 933         if user_can_seek_ns_from_origin_ret_val is not None:
 934
 935             def user_can_seek_ns_from_origin(self, ns_from_origin):
 936                 nonlocal received_ns_from_origin
 937                 received_ns_from_origin = ns_from_origin
 938                 return user_can_seek_ns_from_origin_ret_val
 939
 940         else:
 941             user_can_seek_ns_from_origin = None
 942
 943         if user_seek_ns_from_origin_provided:
 944
 945             def user_seek_ns_from_origin(self, ns_from_origin):
 946                 pass
 947
 948         else:
 949             user_seek_ns_from_origin = None
 950
 951         if iter_can_seek_beginning:
 952
 953             def user_seek_beginning(self):
 954                 pass
 955
 956         else:
 957             user_seek_beginning = None
 958
 959         graph = _setup_seek_test(
 960             MySink,
 961             user_can_seek_ns_from_origin=user_can_seek_ns_from_origin,
 962             user_seek_ns_from_origin=user_seek_ns_from_origin,
 963             user_seek_beginning=user_seek_beginning,
 964             can_seek_forward=iter_can_seek_forward,
 965         )
 966
 967         passed_ns_from_origin = 77
 968         received_ns_from_origin = None
 969         can_seek_ns_from_origin = None
 970         graph.run_once()
 971         self.assertIs(can_seek_ns_from_origin, expected_outcome)
 972
 973         if user_can_seek_ns_from_origin_ret_val is not None:
 974             self.assertEqual(received_ns_from_origin, passed_ns_from_origin)
 975
 976     def test_can_seek_ns_from_origin_user_error(self):
 977         class MySink(bt2._UserSinkComponent):
 978             def __init__(self, config, params, obj):
 979                 self._add_input_port('in')
 980
 981             def _user_graph_is_configured(self):
 982                 self._msg_iter = self._create_message_iterator(self._input_ports['in'])
 983
 984             def _user_consume(self):
 985                 # This is expected to raise.
 986                 self._msg_iter.can_seek_ns_from_origin(2)
 987
 988         def _user_can_seek_ns_from_origin(self, ns_from_origin):
 989             raise ValueError('Joutel')
 990
 991         graph = _setup_seek_test(
 992             MySink,
 993             user_can_seek_ns_from_origin=_user_can_seek_ns_from_origin,
 994             user_seek_ns_from_origin=lambda: None,
 995         )
 996
 997         with self.assertRaises(bt2._Error) as ctx:
 998             graph.run_once()
 999
1000         cause = ctx.exception[0]
1001         self.assertIn('ValueError: Joutel', cause.message)
1002
1003     def test_can_seek_ns_from_origin_wrong_return_value(self):
1004         class MySink(bt2._UserSinkComponent):
1005             def __init__(self, config, params, obj):
1006                 self._add_input_port('in')
1007
1008             def _user_graph_is_configured(self):
1009                 self._msg_iter = self._create_message_iterator(self._input_ports['in'])
1010
1011             def _user_consume(self):
1012                 # This is expected to raise.
1013                 self._msg_iter.can_seek_ns_from_origin(2)
1014
1015         def _user_can_seek_ns_from_origin(self, ns_from_origin):
1016             return 'Nitchequon'
1017
1018         graph = _setup_seek_test(
1019             MySink,
1020             user_can_seek_ns_from_origin=_user_can_seek_ns_from_origin,
1021             user_seek_ns_from_origin=lambda: None,
1022         )
1023
1024         with self.assertRaises(bt2._Error) as ctx:
1025             graph.run_once()
1026
1027         cause = ctx.exception[0]
1028         self.assertIn("TypeError: 'str' is not a 'bool' object", cause.message)
1029
1030     def test_seek_ns_from_origin(self):
1031         class MySink(bt2._UserSinkComponent):
1032             def __init__(self, config, params, obj):
1033                 self._add_input_port('in')
1034
1035             def _user_graph_is_configured(self):
1036                 self._msg_iter = self._create_message_iterator(self._input_ports['in'])
1037
1038             def _user_consume(self):
1039                 self._msg_iter.seek_ns_from_origin(17)
1040
1041         def _user_seek_ns_from_origin(self, ns_from_origin):
1042             nonlocal actual_ns_from_origin
1043             actual_ns_from_origin = ns_from_origin
1044
1045         graph = _setup_seek_test(
1046             MySink, user_seek_ns_from_origin=_user_seek_ns_from_origin
1047         )
1048
1049         actual_ns_from_origin = None
1050         graph.run_once()
1051         self.assertEqual(actual_ns_from_origin, 17)
1052
1053
1054 if __name__ == '__main__':
1055     unittest.main()