Fix: bt2: fix reference counting of messages returned by Python components

[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 9c8277b7f6b1ef2e3e0054d9f51c0fffdd7aae0b..98ec64aba7c5796bff0052a793e1bca60c2be013 100644 (file)
--- a/tests/bindings/python/bt2/test_message_iterator.py
+++ b/tests/bindings/python/bt2/test_message_iterator.py
@@ -1,3 +1,22 @@
+
+#
+# Copyright (C) 2019 EfficiOS Inc.
+#
+# This program is free software; you can redistribute it and/or
+# modify it under the terms of the GNU General Public License
+# as published by the Free Software Foundation; only version 2
+# of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, write to the Free Software
+# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
+#
+
 from bt2 import value
 import collections
 import unittest
@@ -100,6 +119,53 @@ class UserMessageIteratorTestCase(unittest.TestCase):
         self.assertIsNotNone(addr)
         self.assertNotEqual(addr, 0)
 
+    # Test that messages returned by _UserMessageIterator.__next__ remain valid
+    # and can be re-used.
+    def test_reuse_message(self):
+        class MyIter(bt2._UserMessageIterator):
+            def __init__(self, port):
+                tc, sc, ec = port.user_data
+                trace = tc()
+                stream = trace.create_stream(sc)
+                packet = stream.create_packet()
+
+                # This message will be returned twice by __next__.
+                event_message = self._create_event_message(ec, packet)
+
+                self._msgs = [
+                    self._create_stream_beginning_message(stream),
+                    self._create_stream_activity_beginning_message(stream),
+                    self._create_packet_beginning_message(packet),
+                    event_message,
+                    event_message,
+                ]
+
+            def __next__(self):
+                return self._msgs.pop(0)
+
+        class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter):
+            def __init__(self, params):
+                tc = self._create_trace_class()
+                sc = tc.create_stream_class()
+                ec = sc.create_event_class()
+                self._add_output_port('out', (tc, sc, ec))
+
+        graph = bt2.Graph()
+        src = graph.add_component(MySource, 'src')
+        it = graph.create_output_port_message_iterator(src.output_ports['out'])
+
+        # Skip beginning messages.
+        next(it)
+        next(it)
+        next(it)
+
+        msg_ev1 = next(it)
+        msg_ev2 = next(it)
+
+        self.assertIsInstance(msg_ev1, bt2.message._EventMessage)
+        self.assertIsInstance(msg_ev2, bt2.message._EventMessage)
+        self.assertEqual(msg_ev1.addr, msg_ev2.addr)
+
 
 class OutputPortMessageIteratorTestCase(unittest.TestCase):
     def test_component(self):
@@ -167,6 +233,6 @@ class OutputPortMessageIteratorTestCase(unittest.TestCase):
                 self.assertIsInstance(msg, bt2.message._StreamEndMessage)
             else:
                 self.assertIsInstance(msg, bt2.message._EventMessage)
-                self.assertEqual(msg.event.event_class.name, 'salut')
+                self.assertEqual(msg.event.cls.name, 'salut')
                 field = msg.event.payload_field['my_int']
                 self.assertEqual(field, at * 3)