lukecwik commented on a change in pull request #10892: [BEAM-8335] Make
TestStream to/from runner_api include the output_tags property.
URL: https://github.com/apache/beam/pull/10892#discussion_r384689184
##########
File path: sdks/python/apache_beam/runners/direct/test_stream_impl.py
##########
@@ -45,11 +46,69 @@ class _WatermarkController(PTransform):
- If the instance receives an ElementEvent, it emits all specified elements
to the Global Window with the event time set to the element's timestamp.
"""
+ def __init__(self, output_tag):
+ self.output_tag = output_tag
+
def get_windowing(self, _):
return core.Windowing(window.GlobalWindows())
def expand(self, pcoll):
- return pvalue.PCollection.from_(pcoll)
+ ret = pvalue.PCollection.from_(pcoll)
+ ret.tag = self.output_tag
+ return ret
+
+
+class _ExpandableTestStream(PTransform):
+ def __init__(self, test_stream):
+ self.test_stream = test_stream
+
+ def expand(self, pbegin):
+ """Expands the TestStream into the DirectRunner implementation.
+
Review comment:
```suggestion
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
