X-Git-Url: http://git.efficios.com/?a=blobdiff_plain;f=tests%2Fbindings%2Fpython%2Fbt2%2Futils.py;h=7707a531d1a7576c9d32328dda43cf2ed831eea1;hb=f3c9a159782f70dbd0e5dedb37e4a1ef8a6d304e;hp=768462ec09993e7acd6de544c97879e163e1e9ff;hpb=fbbe93021d7cd3793003911931bb3abd5e69596a;p=babeltrace.git

diff --git a/tests/bindings/python/bt2/utils.py b/tests/bindings/python/bt2/utils.py
index 768462ec..7707a531 100644
--- a/tests/bindings/python/bt2/utils.py
+++ b/tests/bindings/python/bt2/utils.py
@@ -1,17 +1,35 @@
+#
+# 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.
+#
+
 import bt2
+import collections.abc
 
 # Run callable `func` in the context of a component's __init__ method.  The
 # callable is passed the Component being instantiated.
 #
 # The value returned by the callable is returned by run_in_component_init.
-
 def run_in_component_init(func):
     class MySink(bt2._UserSinkComponent):
-        def __init__(self, params):
+        def __init__(self, params, obj):
             nonlocal res_bound
             res_bound = func(self)
 
-        def _consume(self):
+        def _user_consume(self):
             pass
 
     g = bt2.Graph()
@@ -31,10 +49,58 @@ def run_in_component_init(func):
     del res_bound
     return res
 
-# Create an empty trace class with default values.
 
+# Create an empty trace class with default values.
 def get_default_trace_class():
     def f(comp_self):
         return comp_self._create_trace_class()
 
     return run_in_component_init(f)
+
+
+# Proxy sink component class.
+#
+# This sink accepts a list of a single item as its initialization
+# object. This sink creates a single input port `in`. When it consumes
+# from this port, it puts the returned message in the initialization
+# list as the first item.
+class TestProxySink(bt2._UserSinkComponent):
+    def __init__(self, params, msg_list):
+        assert msg_list is not None
+        self._msg_list = msg_list
+        self._add_input_port('in')
+
+    def _user_graph_is_configured(self):
+        self._msg_iter = self._create_input_port_message_iterator(
+            self._input_ports['in']
+        )
+
+    def _user_consume(self):
+        assert self._msg_list[0] is None
+        self._msg_list[0] = next(self._msg_iter)
+
+
+# This is a helper message iterator for tests.
+#
+# The constructor accepts a graph and an output port.
+#
+# Internally, it adds a proxy sink to the graph and connects the
+# received output port to the proxy sink's input port. Its __next__()
+# method then uses the proxy sink to transfer the consumed message to
+# the output port message iterator's user.
+#
+# This message iterator cannot seek.
+class TestOutputPortMessageIterator(collections.abc.Iterator):
+    def __init__(self, graph, output_port):
+        self._graph = graph
+        self._msg_list = [None]
+        sink = graph.add_component(TestProxySink, 'test-proxy-sink', obj=self._msg_list)
+        graph.connect_ports(output_port, sink.input_ports['in'])
+
+    def __next__(self):
+        assert self._msg_list[0] is None
+        self._graph.run_once()
+        msg = self._msg_list[0]
+        assert msg is not None
+        self._msg_list[0] = None
+        return msg