bt2: allow getting self output port from user message iterator

[babeltrace.git] / tests / bindings / python / bt2 / test_message_iterator.py

diff --git a/tests/bindings/python/bt2/test_message_iterator.py b/tests/bindings/python/bt2/test_message_iterator.py
index c317ae432f7764a7bfdea06f329d422341c163d2..acd091f872bdca240753a8bd26f64d1e2df4b612 100644 (file)
--- a/tests/bindings/python/bt2/test_message_iterator.py
+++ b/tests/bindings/python/bt2/test_message_iterator.py
@@ -20,6 +20,7 @@ import unittest
 import bt2
 import sys
 from utils import TestOutputPortMessageIterator
+from bt2 import port as bt2_port
 
 
 class UserMessageIteratorTestCase(unittest.TestCase):
@@ -186,6 +187,25 @@ class UserMessageIteratorTestCase(unittest.TestCase):
         graph.run()
         self.assertEqual(salut, 23)
 
+    def test_port(self):
+        class MyIter(bt2._UserMessageIterator):
+            def __init__(self_iter, self_port_output):
+                nonlocal called
+                called = True
+                port = self_iter._port
+                self.assertIs(type(self_port_output), bt2_port._UserComponentOutputPort)
+                self.assertIs(type(port), bt2_port._UserComponentOutputPort)
+                self.assertEqual(self_port_output.addr, port.addr)
+
+        class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
+            def __init__(self, params, obj):
+                self._add_output_port('out')
+
+        called = False
+        graph = self._create_graph(MySource)
+        graph.run()
+        self.assertTrue(called)
+
     def test_addr(self):
         class MyIter(bt2._UserMessageIterator):
             def __init__(self, self_port_output):
@@ -238,15 +258,15 @@ class UserMessageIteratorTestCase(unittest.TestCase):
 
         # Skip beginning messages.
         msg = next(it)
-        self.assertIsInstance(msg, bt2._StreamBeginningMessage)
+        self.assertIs(type(msg), bt2._StreamBeginningMessageConst)
         msg = next(it)
-        self.assertIsInstance(msg, bt2._PacketBeginningMessage)
+        self.assertIs(type(msg), bt2._PacketBeginningMessageConst)
 
         msg_ev1 = next(it)
         msg_ev2 = next(it)
 
-        self.assertIsInstance(msg_ev1, bt2._EventMessage)
-        self.assertIsInstance(msg_ev2, bt2._EventMessage)
+        self.assertIs(type(msg_ev1), bt2._EventMessageConst)
+        self.assertIs(type(msg_ev2), bt2._EventMessageConst)
         self.assertEqual(msg_ev1.addr, msg_ev2.addr)
 
     @staticmethod
@@ -394,14 +414,14 @@ class UserMessageIteratorTestCase(unittest.TestCase):
         msg = None
         MySourceIter, graph = self._setup_seek_beginning_test(MySink)
         graph.run_once()
-        self.assertIsInstance(msg, bt2._StreamBeginningMessage)
+        self.assertIs(type(msg), bt2._StreamBeginningMessageConst)
         graph.run_once()
-        self.assertIsInstance(msg, bt2._PacketBeginningMessage)
+        self.assertIs(type(msg), bt2._PacketBeginningMessageConst)
         do_seek_beginning = True
         graph.run_once()
         do_seek_beginning = False
         graph.run_once()
-        self.assertIsInstance(msg, bt2._StreamBeginningMessage)
+        self.assertIs(type(msg), bt2._StreamBeginningMessageConst)
 
     def test_seek_beginning_user_error(self):
         class MySink(bt2._UserSinkComponent):