X-Git-Url: http://git.efficios.com/?p=babeltrace.git;a=blobdiff_plain;f=tests%2Fbindings%2Fpython%2Fbt2%2Ftest_message.py;h=ef8e6b800a270eef25ab1c15ec4d09532473fcfe;hp=78e8dec066592aae1ad15ffa293d28ba9e240203;hb=6c373cc905e907ecbad698fee38db1d47a981b14;hpb=c1859f693ac2a28d96e7ad492f8d66f74f33686f

diff --git a/tests/bindings/python/bt2/test_message.py b/tests/bindings/python/bt2/test_message.py
index 78e8dec0..ef8e6b80 100644
--- a/tests/bindings/python/bt2/test_message.py
+++ b/tests/bindings/python/bt2/test_message.py
@@ -19,6 +19,7 @@
 import collections
 import unittest
 import bt2
+from utils import TestOutputPortMessageIterator
 
 
 class AllMessagesTestCase(unittest.TestCase):
@@ -148,8 +149,8 @@ class AllMessagesTestCase(unittest.TestCase):
     def test_all_msg_with_cc(self):
         params = {'with_cc': True}
         self._src_comp = self._graph.add_component(self._src, 'my_source', params)
-        self._msg_iter = self._graph.create_output_port_message_iterator(
-            self._src_comp.output_ports['out']
+        self._msg_iter = TestOutputPortMessageIterator(
+            self._graph, self._src_comp.output_ports['out']
         )
 
         for i, msg in enumerate(self._msg_iter):
@@ -204,8 +205,8 @@ class AllMessagesTestCase(unittest.TestCase):
     def test_all_msg_without_cc(self):
         params = {'with_cc': False}
         self._src_comp = self._graph.add_component(self._src, 'my_source', params)
-        self._msg_iter = self._graph.create_output_port_message_iterator(
-            self._src_comp.output_ports['out']
+        self._msg_iter = TestOutputPortMessageIterator(
+            self._graph, self._src_comp.output_ports['out']
         )
 
         for i, msg in enumerate(self._msg_iter):
@@ -273,8 +274,8 @@ class AllMessagesTestCase(unittest.TestCase):
         params = {'with_cc': True, 'with_stream_msgs_clock_snapshots': True}
 
         self._src_comp = self._graph.add_component(self._src, 'my_source', params)
-        self._msg_iter = self._graph.create_output_port_message_iterator(
-            self._src_comp.output_ports['out']
+        self._msg_iter = TestOutputPortMessageIterator(
+            self._graph, self._src_comp.output_ports['out']
         )
         msgs = list(self._msg_iter)