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