X-Git-Url: http://git.efficios.com/?a=blobdiff_plain;f=tests%2Fbindings%2Fpython%2Fbt2%2Ftest_message_iterator.py;h=40adfa60d847fc16e652540b9811114d23d3f833;hb=f0a42b33ac3951cd5cb2ee0f66ac04437a681621;hp=3cd27e5b77325536a7c8735756edd539bf76ec01;hpb=6c373cc905e907ecbad698fee38db1d47a981b14;p=babeltrace.git

diff --git a/tests/bindings/python/bt2/test_message_iterator.py b/tests/bindings/python/bt2/test_message_iterator.py
index 3cd27e5b..40adfa60 100644
--- a/tests/bindings/python/bt2/test_message_iterator.py
+++ b/tests/bindings/python/bt2/test_message_iterator.py
@@ -16,11 +16,9 @@
 # Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
 #
 
-from bt2 import value
-import collections
 import unittest
-import copy
 import bt2
+import sys
 from utils import TestOutputPortMessageIterator
 
 
@@ -117,6 +115,45 @@ class UserMessageIteratorTestCase(unittest.TestCase):
         self.assertTrue(src_iter_initialized)
         self.assertTrue(flt_iter_initialized)
 
+    def test_create_user_error(self):
+        # This tests both error handling by
+        # _UserSinkComponent._create_input_port_message_iterator
+        # and _UserMessageIterator._create_input_port_message_iterator, as they
+        # are both used in the graph.
+        class MySourceIter(bt2._UserMessageIterator):
+            def __init__(self, self_port_output):
+                raise ValueError('Very bad error')
+
+        class MySource(bt2._UserSourceComponent, message_iterator_class=MySourceIter):
+            def __init__(self, params, obj):
+                self._add_output_port('out')
+
+        class MyFilterIter(bt2._UserMessageIterator):
+            def __init__(self, self_port_output):
+                # This is expected to raise because of the error in
+                # MySourceIter.__init__.
+                self._create_input_port_message_iterator(
+                    self._component._input_ports['in']
+                )
+
+        class MyFilter(bt2._UserFilterComponent, message_iterator_class=MyFilterIter):
+            def __init__(self, params, obj):
+                self._add_input_port('in')
+                self._add_output_port('out')
+
+        graph = self._create_graph(MySource, MyFilter)
+
+        with self.assertRaises(bt2._Error) as ctx:
+            graph.run()
+
+        exc = ctx.exception
+        cause = exc[0]
+
+        self.assertIsInstance(cause, bt2._MessageIteratorErrorCause)
+        self.assertEqual(cause.component_name, 'src')
+        self.assertEqual(cause.component_output_port_name, 'out')
+        self.assertIn('ValueError: Very bad error', cause.message)
+
     def test_finalize(self):
         class MyIter(bt2._UserMessageIterator):
             def _user_finalize(self):
@@ -201,15 +238,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
@@ -357,14 +394,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):
@@ -405,10 +442,12 @@ class UserMessageIteratorTestCase(unittest.TestCase):
         src = graph.add_component(MySource, 'src')
         it = TestOutputPortMessageIterator(graph, src.output_ports['out'])
 
-        # The initial refcount of Py_None was in the 7000, so 100000 iterations
-        # should be enough to catch the bug even if there are small differences
+        # Three times the initial ref count of `None` iterations should
+        # be enough to catch the bug even if there are small differences
         # between configurations.
-        for i in range(100000):
+        none_ref_count = sys.getrefcount(None) * 3
+
+        for i in range(none_ref_count):
             with self.assertRaises(bt2.TryAgain):
                 next(it)