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
 383 def _setup_seek_test(
 384     sink_cls,
 385     user_seek_beginning=None,
 386     user_can_seek_beginning=None,
 387     user_seek_ns_from_origin=None,
 388     user_can_seek_ns_from_origin=None,
 389     can_seek_forward=False,
 390 ):
 391     class MySourceIter(bt2._UserMessageIterator):
 392         def __init__(self, config, port):
 393             tc, sc, ec = port.user_data
 394             trace = tc()
 395             stream = trace.create_stream(sc)
 396             packet = stream.create_packet()
 397
 398             self._msgs = [
 399                 self._create_stream_beginning_message(stream),
 400                 self._create_packet_beginning_message(packet),
 401                 self._create_event_message(ec, packet),
 402                 self._create_event_message(ec, packet),
 403                 self._create_packet_end_message(packet),
 404                 self._create_stream_end_message(stream),
 405             ]
 406             self._at = 0
 407             config.can_seek_forward = can_seek_forward
 408
 409         def __next__(self):
 410             if self._at < len(self._msgs):
 411                 msg = self._msgs[self._at]
 412                 self._at += 1
 413                 return msg
 414             else:
 415                 raise StopIteration
 416
 417     if user_seek_beginning is not None:
 418         MySourceIter._user_seek_beginning = user_seek_beginning
 419
 420     if user_can_seek_beginning is not None:
 421         MySourceIter._user_can_seek_beginning = user_can_seek_beginning
 422
 423     if user_seek_ns_from_origin is not None:
 424         MySourceIter._user_seek_ns_from_origin = user_seek_ns_from_origin
 425
 426     if user_can_seek_ns_from_origin is not None:
 427         MySourceIter._user_can_seek_ns_from_origin = user_can_seek_ns_from_origin
 428
 429     class MySource(bt2._UserSourceComponent, message_iterator_class=MySourceIter):
 430         def __init__(self, config, params, obj):
 431             tc = self._create_trace_class()
 432             sc = tc.create_stream_class(supports_packets=True)
 433             ec = sc.create_event_class()
 434
 435             self._add_output_port('out', (tc, sc, ec))
 436
 437     class MyFilterIter(bt2._UserMessageIterator):
 438         def __init__(self, config, port):
 439             self._upstream_iter = self._create_input_port_message_iterator(
 440                 self._component._input_ports['in']
 441             )
 442             config.can_seek_forward = self._upstream_iter.can_seek_forward
 443
 444         def __next__(self):
 445             return next(self._upstream_iter)
 446
 447         def _user_can_seek_beginning(self):
 448             return self._upstream_iter.can_seek_beginning()
 449
 450         def _user_seek_beginning(self):
 451             self._upstream_iter.seek_beginning()
 452
 453         def _user_can_seek_ns_from_origin(self, ns_from_origin):
 454             return self._upstream_iter.can_seek_ns_from_origin(ns_from_origin)
 455
 456         def _user_seek_ns_from_origin(self, ns_from_origin):
 457             self._upstream_iter.seek_ns_from_origin(ns_from_origin)
 458
 459     class MyFilter(bt2._UserFilterComponent, message_iterator_class=MyFilterIter):
 460         def __init__(self, config, params, obj):
 461             self._add_input_port('in')
 462             self._add_output_port('out')
 463
 464     return _create_graph(MySource, sink_cls, flt_comp_cls=MyFilter)
 465
 466
 467 class UserMessageIteratorSeekBeginningTestCase(unittest.TestCase):
 468     def test_can_seek_beginning(self):
 469         class MySink(bt2._UserSinkComponent):
 470             def __init__(self, config, params, obj):
 471                 self._add_input_port('in')
 472
 473             def _user_graph_is_configured(self):
 474                 self._msg_iter = self._create_input_port_message_iterator(
 475                     self._input_ports['in']
 476                 )
 477
 478             def _user_consume(self):
 479                 nonlocal can_seek_beginning
 480                 can_seek_beginning = self._msg_iter.can_seek_beginning()
 481
 482         def _user_can_seek_beginning(self):
 483             nonlocal input_port_iter_can_seek_beginning
 484             return input_port_iter_can_seek_beginning
 485
 486         graph = _setup_seek_test(
 487             MySink, user_can_seek_beginning=_user_can_seek_beginning
 488         )
 489
 490         input_port_iter_can_seek_beginning = True
 491         can_seek_beginning = None
 492         graph.run_once()
 493         self.assertIs(can_seek_beginning, True)
 494
 495         input_port_iter_can_seek_beginning = False
 496         can_seek_beginning = None
 497         graph.run_once()
 498         self.assertIs(can_seek_beginning, False)
 499
 500     def test_no_can_seek_beginning_with_seek_beginning(self):
 501         # Test an iterator without a _user_can_seek_beginning method, but with
 502         # a _user_seek_beginning method.
 503         class MySink(bt2._UserSinkComponent):
 504             def __init__(self, config, params, obj):
 505                 self._add_input_port('in')
 506
 507             def _user_graph_is_configured(self):
 508                 self._msg_iter = self._create_input_port_message_iterator(
 509                     self._input_ports['in']
 510                 )
 511
 512             def _user_consume(self):
 513                 nonlocal can_seek_beginning
 514                 can_seek_beginning = self._msg_iter.can_seek_beginning()
 515
 516         def _user_seek_beginning(self):
 517             pass
 518
 519         graph = _setup_seek_test(MySink, user_seek_beginning=_user_seek_beginning)
 520         can_seek_beginning = None
 521         graph.run_once()
 522         self.assertIs(can_seek_beginning, True)
 523
 524     def test_no_can_seek_beginning(self):
 525         # Test an iterator without a _user_can_seek_beginning method, without
 526         # a _user_seek_beginning method.
 527         class MySink(bt2._UserSinkComponent):
 528             def __init__(self, config, params, obj):
 529                 self._add_input_port('in')
 530
 531             def _user_graph_is_configured(self):
 532                 self._msg_iter = self._create_input_port_message_iterator(
 533                     self._input_ports['in']
 534                 )
 535
 536             def _user_consume(self):
 537                 nonlocal can_seek_beginning
 538                 can_seek_beginning = self._msg_iter.can_seek_beginning()
 539
 540         graph = _setup_seek_test(MySink)
 541         can_seek_beginning = None
 542         graph.run_once()
 543         self.assertIs(can_seek_beginning, False)
 544
 545     def test_can_seek_beginning_user_error(self):
 546         class MySink(bt2._UserSinkComponent):
 547             def __init__(self, config, params, obj):
 548                 self._add_input_port('in')
 549
 550             def _user_graph_is_configured(self):
 551                 self._msg_iter = self._create_input_port_message_iterator(
 552                     self._input_ports['in']
 553                 )
 554
 555             def _user_consume(self):
 556                 # This is expected to raise.
 557                 self._msg_iter.can_seek_beginning()
 558
 559         def _user_can_seek_beginning(self):
 560             raise ValueError('moustiquaire')
 561
 562         graph = _setup_seek_test(
 563             MySink, user_can_seek_beginning=_user_can_seek_beginning
 564         )
 565
 566         with self.assertRaises(bt2._Error) as ctx:
 567             graph.run_once()
 568
 569         cause = ctx.exception[0]
 570         self.assertIn('ValueError: moustiquaire', cause.message)
 571
 572     def test_can_seek_beginning_wrong_return_value(self):
 573         class MySink(bt2._UserSinkComponent):
 574             def __init__(self, config, params, obj):
 575                 self._add_input_port('in')
 576
 577             def _user_graph_is_configured(self):
 578                 self._msg_iter = self._create_input_port_message_iterator(
 579                     self._input_ports['in']
 580                 )
 581
 582             def _user_consume(self):
 583                 # This is expected to raise.
 584                 self._msg_iter.can_seek_beginning()
 585
 586         def _user_can_seek_beginning(self):
 587             return 'Amqui'
 588
 589         graph = _setup_seek_test(
 590             MySink, user_can_seek_beginning=_user_can_seek_beginning
 591         )
 592
 593         with self.assertRaises(bt2._Error) as ctx:
 594             graph.run_once()
 595
 596         cause = ctx.exception[0]
 597         self.assertIn("TypeError: 'str' is not a 'bool' object", cause.message)
 598
 599     def test_seek_beginning(self):
 600         class MySink(bt2._UserSinkComponent):
 601             def __init__(self, config, params, obj):
 602                 self._add_input_port('in')
 603
 604             def _user_graph_is_configured(self):
 605                 self._msg_iter = self._create_input_port_message_iterator(
 606                     self._input_ports['in']
 607                 )
 608
 609             def _user_consume(self):
 610                 nonlocal do_seek_beginning
 611                 nonlocal msg
 612
 613                 if do_seek_beginning:
 614                     self._msg_iter.seek_beginning()
 615                     return
 616
 617                 msg = next(self._msg_iter)
 618
 619         def _user_seek_beginning(self):
 620             self._at = 0
 621
 622         msg = None
 623         graph = _setup_seek_test(MySink, user_seek_beginning=_user_seek_beginning)
 624
 625         # Consume message.
 626         do_seek_beginning = False
 627         graph.run_once()
 628         self.assertIs(type(msg), bt2._StreamBeginningMessageConst)
 629
 630         # Consume message.
 631         graph.run_once()
 632         self.assertIs(type(msg), bt2._PacketBeginningMessageConst)
 633
 634         # Seek beginning.
 635         do_seek_beginning = True
 636         graph.run_once()
 637
 638         # Consume message.
 639         do_seek_beginning = False
 640         graph.run_once()
 641         self.assertIs(type(msg), bt2._StreamBeginningMessageConst)
 642
 643     def test_seek_beginning_user_error(self):
 644         class MySink(bt2._UserSinkComponent):
 645             def __init__(self, config, params, obj):
 646                 self._add_input_port('in')
 647
 648             def _user_graph_is_configured(self):
 649                 self._msg_iter = self._create_input_port_message_iterator(
 650                     self._input_ports['in']
 651                 )
 652
 653             def _user_consume(self):
 654                 self._msg_iter.seek_beginning()
 655
 656         def _user_seek_beginning(self):
 657             raise ValueError('ouch')
 658
 659         graph = _setup_seek_test(MySink, user_seek_beginning=_user_seek_beginning)
 660
 661         with self.assertRaises(bt2._Error):
 662             graph.run_once()
 663
 664
 665 class UserMessageIteratorSeekNsFromOriginTestCase(unittest.TestCase):
 666     def test_can_seek_ns_from_origin_returns_true(self):
 667         # Test the case where:
 668         #
 669         #   - can_seek_ns_from_origin: returns True
 670         #   - seek_ns_from_origin provided: Don't care
 671         #   - can the iterator seek beginning: Don't care
 672         #   - can the iterator seek forward: Don't care
 673         #
 674         # We expect iter.can_seek_ns_from_origin to return True.
 675         for user_seek_ns_from_origin_provided in (False, True):
 676             for iter_can_seek_beginning in (False, True):
 677                 for iter_can_seek_forward in (False, True):
 678                     self._can_seek_ns_from_origin_test(
 679                         expected_outcome=True,
 680                         user_can_seek_ns_from_origin_ret_val=True,
 681                         user_seek_ns_from_origin_provided=user_seek_ns_from_origin_provided,
 682                         iter_can_seek_beginning=iter_can_seek_beginning,
 683                         iter_can_seek_forward=iter_can_seek_forward,
 684                     )
 685
 686     def test_can_seek_ns_from_origin_returns_false_can_seek_beginning_forward_seekable(
 687         self
 688     ):
 689         # Test the case where:
 690         #
 691         #   - can_seek_ns_from_origin: returns False
 692         #   - seek_ns_from_origin provided: Don't care
 693         #   - can the iterator seek beginning: Yes
 694         #   - can the iterator seek forward: Yes
 695         #
 696         # We expect iter.can_seek_ns_from_origin to return True.
 697         for user_seek_ns_from_origin_provided in (False, True):
 698             self._can_seek_ns_from_origin_test(
 699                 expected_outcome=True,
 700                 user_can_seek_ns_from_origin_ret_val=False,
 701                 user_seek_ns_from_origin_provided=user_seek_ns_from_origin_provided,
 702                 iter_can_seek_beginning=True,
 703                 iter_can_seek_forward=True,
 704             )
 705
 706     def test_can_seek_ns_from_origin_returns_false_can_seek_beginning_not_forward_seekable(
 707         self
 708     ):
 709         # Test the case where:
 710         #
 711         #   - can_seek_ns_from_origin: returns False
 712         #   - seek_ns_from_origin provided: Don't care
 713         #   - can the iterator seek beginning: Yes
 714         #   - can the iterator seek forward: No
 715         #
 716         # We expect iter.can_seek_ns_from_origin to return False.
 717         for user_seek_ns_from_origin_provided in (False, True):
 718             self._can_seek_ns_from_origin_test(
 719                 expected_outcome=False,
 720                 user_can_seek_ns_from_origin_ret_val=False,
 721                 user_seek_ns_from_origin_provided=user_seek_ns_from_origin_provided,
 722                 iter_can_seek_beginning=True,
 723                 iter_can_seek_forward=False,
 724             )
 725
 726     def test_can_seek_ns_from_origin_returns_false_cant_seek_beginning_forward_seekable(
 727         self
 728     ):
 729         # Test the case where:
 730         #
 731         #   - can_seek_ns_from_origin: returns False
 732         #   - seek_ns_from_origin provided: Don't care
 733         #   - can the iterator seek beginning: No
 734         #   - can the iterator seek forward: Yes
 735         #
 736         # We expect iter.can_seek_ns_from_origin to return False.
 737         # for user_seek_ns_from_origin_provided in (False, True):
 738         self._can_seek_ns_from_origin_test(
 739             expected_outcome=False,
 740             user_can_seek_ns_from_origin_ret_val=False,
 741             user_seek_ns_from_origin_provided=False,
 742             iter_can_seek_beginning=False,
 743             iter_can_seek_forward=True,
 744         )
 745
 746     def test_can_seek_ns_from_origin_returns_false_cant_seek_beginning_not_forward_seekable(
 747         self
 748     ):
 749         # Test the case where:
 750         #
 751         #   - can_seek_ns_from_origin: returns False
 752         #   - seek_ns_from_origin provided: Don't care
 753         #   - can the iterator seek beginning: No
 754         #   - can the iterator seek forward: No
 755         #
 756         # We expect iter.can_seek_ns_from_origin to return False.
 757         for user_seek_ns_from_origin_provided in (False, True):
 758             self._can_seek_ns_from_origin_test(
 759                 expected_outcome=False,
 760                 user_can_seek_ns_from_origin_ret_val=False,
 761                 user_seek_ns_from_origin_provided=user_seek_ns_from_origin_provided,
 762                 iter_can_seek_beginning=False,
 763                 iter_can_seek_forward=False,
 764             )
 765
 766     def test_no_can_seek_ns_from_origin_seek_ns_from_origin(self):
 767         # Test the case where:
 768         #
 769         #   - can_seek_ns_from_origin: Not provided
 770         #   - seek_ns_from_origin provided: Yes
 771         #   - can the iterator seek beginning: Don't care
 772         #   - can the iterator seek forward: Don't care
 773         #
 774         # We expect iter.can_seek_ns_from_origin to return True.
 775         for iter_can_seek_beginning in (False, True):
 776             for iter_can_seek_forward in (False, True):
 777                 self._can_seek_ns_from_origin_test(
 778                     expected_outcome=True,
 779                     user_can_seek_ns_from_origin_ret_val=None,
 780                     user_seek_ns_from_origin_provided=True,
 781                     iter_can_seek_beginning=iter_can_seek_beginning,
 782                     iter_can_seek_forward=iter_can_seek_forward,
 783                 )
 784
 785     def test_no_can_seek_ns_from_origin_no_seek_ns_from_origin_can_seek_beginning_forward_seekable(
 786         self
 787     ):
 788         # Test the case where:
 789         #
 790         #   - can_seek_ns_from_origin: Not provided
 791         #   - seek_ns_from_origin provided: Not provided
 792         #   - can the iterator seek beginning: Yes
 793         #   - can the iterator seek forward: Yes
 794         #
 795         # We expect iter.can_seek_ns_from_origin to return True.
 796         self._can_seek_ns_from_origin_test(
 797             expected_outcome=True,
 798             user_can_seek_ns_from_origin_ret_val=None,
 799             user_seek_ns_from_origin_provided=False,
 800             iter_can_seek_beginning=True,
 801             iter_can_seek_forward=True,
 802         )
 803
 804     def test_no_can_seek_ns_from_origin_no_seek_ns_from_origin_can_seek_beginning_not_forward_seekable(
 805         self
 806     ):
 807         # Test the case where:
 808         #
 809         #   - can_seek_ns_from_origin: Not provided
 810         #   - seek_ns_from_origin provided: Not provided
 811         #   - can the iterator seek beginning: Yes
 812         #   - can the iterator seek forward: No
 813         #
 814         # We expect iter.can_seek_ns_from_origin to return False.
 815         self._can_seek_ns_from_origin_test(
 816             expected_outcome=False,
 817             user_can_seek_ns_from_origin_ret_val=None,
 818             user_seek_ns_from_origin_provided=False,
 819             iter_can_seek_beginning=True,
 820             iter_can_seek_forward=False,
 821         )
 822
 823     def test_no_can_seek_ns_from_origin_no_seek_ns_from_origin_cant_seek_beginning_forward_seekable(
 824         self
 825     ):
 826         # Test the case where:
 827         #
 828         #   - can_seek_ns_from_origin: Not provided
 829         #   - seek_ns_from_origin provided: Not provided
 830         #   - can the iterator seek beginning: No
 831         #   - can the iterator seek forward: Yes
 832         #
 833         # We expect iter.can_seek_ns_from_origin to return False.
 834         self._can_seek_ns_from_origin_test(
 835             expected_outcome=False,
 836             user_can_seek_ns_from_origin_ret_val=None,
 837             user_seek_ns_from_origin_provided=False,
 838             iter_can_seek_beginning=False,
 839             iter_can_seek_forward=True,
 840         )
 841
 842     def test_no_can_seek_ns_from_origin_no_seek_ns_from_origin_cant_seek_beginning_not_forward_seekable(
 843         self
 844     ):
 845         # Test the case where:
 846         #
 847         #   - can_seek_ns_from_origin: Not provided
 848         #   - seek_ns_from_origin provided: Not provided
 849         #   - can the iterator seek beginning: No
 850         #   - can the iterator seek forward: No
 851         #
 852         # We expect iter.can_seek_ns_from_origin to return False.
 853         self._can_seek_ns_from_origin_test(
 854             expected_outcome=False,
 855             user_can_seek_ns_from_origin_ret_val=None,
 856             user_seek_ns_from_origin_provided=False,
 857             iter_can_seek_beginning=False,
 858             iter_can_seek_forward=False,
 859         )
 860
 861     def _can_seek_ns_from_origin_test(
 862         self,
 863         expected_outcome,
 864         user_can_seek_ns_from_origin_ret_val,
 865         user_seek_ns_from_origin_provided,
 866         iter_can_seek_beginning,
 867         iter_can_seek_forward,
 868     ):
 869         class MySink(bt2._UserSinkComponent):
 870             def __init__(self, config, params, obj):
 871                 self._add_input_port('in')
 872
 873             def _user_graph_is_configured(self):
 874                 self._msg_iter = self._create_input_port_message_iterator(
 875                     self._input_ports['in']
 876                 )
 877
 878             def _user_consume(self):
 879                 nonlocal can_seek_ns_from_origin
 880                 can_seek_ns_from_origin = self._msg_iter.can_seek_ns_from_origin(
 881                     passed_ns_from_origin
 882                 )
 883
 884         if user_can_seek_ns_from_origin_ret_val is not None:
 885
 886             def user_can_seek_ns_from_origin(self, ns_from_origin):
 887                 nonlocal received_ns_from_origin
 888                 received_ns_from_origin = ns_from_origin
 889                 return user_can_seek_ns_from_origin_ret_val
 890
 891         else:
 892             user_can_seek_ns_from_origin = None
 893
 894         if user_seek_ns_from_origin_provided:
 895
 896             def user_seek_ns_from_origin(self, ns_from_origin):
 897                 pass
 898
 899         else:
 900             user_seek_ns_from_origin = None
 901
 902         if iter_can_seek_beginning:
 903
 904             def user_seek_beginning(self):
 905                 pass
 906
 907         else:
 908             user_seek_beginning = None
 909
 910         graph = _setup_seek_test(
 911             MySink,
 912             user_can_seek_ns_from_origin=user_can_seek_ns_from_origin,
 913             user_seek_ns_from_origin=user_seek_ns_from_origin,
 914             user_seek_beginning=user_seek_beginning,
 915             can_seek_forward=iter_can_seek_forward,
 916         )
 917
 918         passed_ns_from_origin = 77
 919         received_ns_from_origin = None
 920         can_seek_ns_from_origin = None
 921         graph.run_once()
 922         self.assertIs(can_seek_ns_from_origin, expected_outcome)
 923
 924         if user_can_seek_ns_from_origin_ret_val is not None:
 925             self.assertEqual(received_ns_from_origin, passed_ns_from_origin)
 926
 927     def test_can_seek_ns_from_origin_user_error(self):
 928         class MySink(bt2._UserSinkComponent):
 929             def __init__(self, config, params, obj):
 930                 self._add_input_port('in')
 931
 932             def _user_graph_is_configured(self):
 933                 self._msg_iter = self._create_input_port_message_iterator(
 934                     self._input_ports['in']
 935                 )
 936
 937             def _user_consume(self):
 938                 # This is expected to raise.
 939                 self._msg_iter.can_seek_ns_from_origin(2)
 940
 941         def _user_can_seek_ns_from_origin(self, ns_from_origin):
 942             raise ValueError('Joutel')
 943
 944         graph = _setup_seek_test(
 945             MySink, user_can_seek_ns_from_origin=_user_can_seek_ns_from_origin
 946         )
 947
 948         with self.assertRaises(bt2._Error) as ctx:
 949             graph.run_once()
 950
 951         cause = ctx.exception[0]
 952         self.assertIn('ValueError: Joutel', cause.message)
 953
 954     def test_can_seek_ns_from_origin_wrong_return_value(self):
 955         class MySink(bt2._UserSinkComponent):
 956             def __init__(self, config, params, obj):
 957                 self._add_input_port('in')
 958
 959             def _user_graph_is_configured(self):
 960                 self._msg_iter = self._create_input_port_message_iterator(
 961                     self._input_ports['in']
 962                 )
 963
 964             def _user_consume(self):
 965                 # This is expected to raise.
 966                 self._msg_iter.can_seek_ns_from_origin(2)
 967
 968         def _user_can_seek_ns_from_origin(self, ns_from_origin):
 969             return 'Nitchequon'
 970
 971         graph = _setup_seek_test(
 972             MySink, user_can_seek_ns_from_origin=_user_can_seek_ns_from_origin
 973         )
 974
 975         with self.assertRaises(bt2._Error) as ctx:
 976             graph.run_once()
 977
 978         cause = ctx.exception[0]
 979         self.assertIn("TypeError: 'str' is not a 'bool' object", cause.message)
 980
 981     def test_seek_ns_from_origin(self):
 982         class MySink(bt2._UserSinkComponent):
 983             def __init__(self, config, params, obj):
 984                 self._add_input_port('in')
 985
 986             def _user_graph_is_configured(self):
 987                 self._msg_iter = self._create_input_port_message_iterator(
 988                     self._input_ports['in']
 989                 )
 990
 991             def _user_consume(self):
 992                 self._msg_iter.seek_ns_from_origin(17)
 993
 994         def _user_seek_ns_from_origin(self, ns_from_origin):
 995             nonlocal actual_ns_from_origin
 996             actual_ns_from_origin = ns_from_origin
 997
 998         graph = _setup_seek_test(
 999             MySink, user_seek_ns_from_origin=_user_seek_ns_from_origin
1000         )
1001
1002         actual_ns_from_origin = None
1003         graph.run_once()
1004         self.assertEqual(actual_ns_from_origin, 17)
1005
1006
1007 if __name__ == '__main__':
1008     unittest.main()