tests/bindings/python/bt2/test_message_iterator.py

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