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

diff --git a/tests/bindings/python/bt2/utils.py b/tests/bindings/python/bt2/utils.py
index fb4003a3..7707a531 100644
--- a/tests/bindings/python/bt2/utils.py
+++ b/tests/bindings/python/bt2/utils.py
@@ -17,23 +17,19 @@
 #
 
 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):
-            pass
-
-        def _graph_is_configured(self):
+        def _user_consume(self):
             pass
 
     g = bt2.Graph()
@@ -55,10 +51,56 @@ def run_in_component_init(func):
 
 
 # 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