| # | |
| # Licensed to the Apache Software Foundation (ASF) under one or more | |
| # contributor license agreements. See the NOTICE file distributed with | |
| # this work for additional information regarding copyright ownership. | |
| # The ASF licenses this file to You under the Apache License, Version 2.0 | |
| # (the "License"); you may not use this file except in compliance with | |
| # the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| """Unit tests for the test_stream module.""" | |
| # pytype: skip-file | |
| import unittest | |
| import apache_beam as beam | |
| from apache_beam.options.pipeline_options import PipelineOptions | |
| from apache_beam.options.pipeline_options import StandardOptions | |
| from apache_beam.options.pipeline_options import TypeOptions | |
| from apache_beam.portability import common_urns | |
| from apache_beam.portability.api import beam_interactive_api_pb2 | |
| from apache_beam.portability.api import beam_runner_api_pb2 | |
| from apache_beam.testing.test_pipeline import TestPipeline | |
| from apache_beam.testing.test_stream import ElementEvent | |
| from apache_beam.testing.test_stream import OutputFormat | |
| from apache_beam.testing.test_stream import ProcessingTimeEvent | |
| from apache_beam.testing.test_stream import ReverseTestStream | |
| from apache_beam.testing.test_stream import TestStream | |
| from apache_beam.testing.test_stream import WatermarkEvent | |
| from apache_beam.testing.test_stream import WindowedValueHolder | |
| from apache_beam.testing.test_stream_service import TestStreamServiceController | |
| from apache_beam.testing.util import assert_that | |
| from apache_beam.testing.util import equal_to | |
| from apache_beam.testing.util import equal_to_per_window | |
| from apache_beam.transforms import trigger | |
| from apache_beam.transforms import window | |
| from apache_beam.transforms.window import FixedWindows | |
| from apache_beam.transforms.window import TimestampedValue | |
| from apache_beam.utils import timestamp | |
| from apache_beam.utils.timestamp import Timestamp | |
| from apache_beam.utils.windowed_value import PaneInfo | |
| from apache_beam.utils.windowed_value import PaneInfoTiming | |
| from apache_beam.utils.windowed_value import WindowedValue | |
| class TestStreamTest(unittest.TestCase): | |
| def test_basic_test_stream(self): | |
| test_stream = (TestStream() | |
| .advance_watermark_to(0) | |
| .add_elements([ | |
| 'a', | |
| WindowedValue('b', 3, []), | |
| TimestampedValue('c', 6)]) | |
| .advance_processing_time(10) | |
| .advance_watermark_to(8) | |
| .add_elements(['d']) | |
| .advance_watermark_to_infinity()) # yapf: disable | |
| self.assertEqual( | |
| test_stream._events, | |
| [ | |
| WatermarkEvent(0), | |
| ElementEvent([ | |
| TimestampedValue('a', 0), | |
| TimestampedValue('b', 3), | |
| TimestampedValue('c', 6), | |
| ]), | |
| ProcessingTimeEvent(10), | |
| WatermarkEvent(8), | |
| ElementEvent([ | |
| TimestampedValue('d', 8), | |
| ]), | |
| WatermarkEvent(timestamp.MAX_TIMESTAMP), | |
| ]) | |
| def test_test_stream_errors(self): | |
| with self.assertRaises( | |
| AssertionError, msg=('Watermark must strictly-monotonically advance.')): | |
| _ = (TestStream().advance_watermark_to(5).advance_watermark_to(4)) | |
| with self.assertRaises( | |
| AssertionError, | |
| msg=('Must advance processing time by positive amount.')): | |
| _ = (TestStream().advance_processing_time(-1)) | |
| with self.assertRaises( | |
| AssertionError, | |
| msg=('Element timestamp must be before timestamp.MAX_TIMESTAMP.')): | |
| _ = ( | |
| TestStream().add_elements( | |
| [TimestampedValue('a', timestamp.MAX_TIMESTAMP)])) | |
| def test_basic_execution(self): | |
| test_stream = (TestStream() | |
| .advance_watermark_to(10) | |
| .add_elements(['a', 'b', 'c']) | |
| .advance_watermark_to(20) | |
| .add_elements(['d']) | |
| .add_elements(['e']) | |
| .advance_processing_time(10) | |
| .advance_watermark_to(300) | |
| .add_elements([TimestampedValue('late', 12)]) | |
| .add_elements([TimestampedValue('last', 310)]) | |
| .advance_watermark_to_infinity()) # yapf: disable | |
| class RecordFn(beam.DoFn): | |
| def process( | |
| self, | |
| element=beam.DoFn.ElementParam, | |
| timestamp=beam.DoFn.TimestampParam): | |
| yield (element, timestamp) | |
| options = PipelineOptions() | |
| options.view_as(StandardOptions).streaming = True | |
| with TestPipeline(options=options) as p: | |
| my_record_fn = RecordFn() | |
| records = p | test_stream | beam.ParDo(my_record_fn) | |
| assert_that( | |
| records, | |
| equal_to([ | |
| ('a', timestamp.Timestamp(10)), | |
| ('b', timestamp.Timestamp(10)), | |
| ('c', timestamp.Timestamp(10)), | |
| ('d', timestamp.Timestamp(20)), | |
| ('e', timestamp.Timestamp(20)), | |
| ('late', timestamp.Timestamp(12)), | |
| ('last', timestamp.Timestamp(310)), | |
| ])) | |
| def test_multiple_outputs(self): | |
| """Tests that the TestStream supports emitting to multiple PCollections.""" | |
| letters_elements = [ | |
| TimestampedValue('a', 6), | |
| TimestampedValue('b', 7), | |
| TimestampedValue('c', 8), | |
| ] | |
| numbers_elements = [ | |
| TimestampedValue('1', 11), | |
| TimestampedValue('2', 12), | |
| TimestampedValue('3', 13), | |
| ] | |
| test_stream = (TestStream() | |
| .advance_watermark_to(5, tag='letters') | |
| .add_elements(letters_elements, tag='letters') | |
| .advance_watermark_to(10, tag='numbers') | |
| .add_elements(numbers_elements, tag='numbers')) # yapf: disable | |
| class RecordFn(beam.DoFn): | |
| def process( | |
| self, | |
| element=beam.DoFn.ElementParam, | |
| timestamp=beam.DoFn.TimestampParam): | |
| yield (element, timestamp) | |
| options = StandardOptions(streaming=True) | |
| p = TestPipeline(options=options) | |
| main = p | test_stream | |
| letters = main['letters'] | 'record letters' >> beam.ParDo(RecordFn()) | |
| numbers = main['numbers'] | 'record numbers' >> beam.ParDo(RecordFn()) | |
| assert_that( | |
| letters, | |
| equal_to([('a', Timestamp(6)), ('b', Timestamp(7)), | |
| ('c', Timestamp(8))]), | |
| label='assert letters') | |
| assert_that( | |
| numbers, | |
| equal_to([('1', Timestamp(11)), ('2', Timestamp(12)), | |
| ('3', Timestamp(13))]), | |
| label='assert numbers') | |
| p.run() | |
| def test_multiple_outputs_with_watermark_advancement(self): | |
| """Tests that the TestStream can independently control output watermarks.""" | |
| # Purposely set the watermark of numbers to 20 then letters to 5 to test | |
| # that the watermark advancement is per PCollection. | |
| # | |
| # This creates two PCollections, (a, b, c) and (1, 2, 3). These will be | |
| # emitted at different times so that they will have different windows. The | |
| # watermark advancement is checked by checking their windows. If the | |
| # watermark does not advance, then the windows will be [-inf, -inf). If the | |
| # windows do not advance separately, then the PCollections will both | |
| # windowed in [15, 30). | |
| letters_elements = [ | |
| TimestampedValue('a', 6), | |
| TimestampedValue('b', 7), | |
| TimestampedValue('c', 8), | |
| ] | |
| numbers_elements = [ | |
| TimestampedValue('1', 21), | |
| TimestampedValue('2', 22), | |
| TimestampedValue('3', 23), | |
| ] | |
| test_stream = (TestStream() | |
| .advance_watermark_to(0, tag='letters') | |
| .advance_watermark_to(0, tag='numbers') | |
| .advance_watermark_to(20, tag='numbers') | |
| .advance_watermark_to(5, tag='letters') | |
| .add_elements(letters_elements, tag='letters') | |
| .advance_watermark_to(10, tag='letters') | |
| .add_elements(numbers_elements, tag='numbers') | |
| .advance_watermark_to(30, tag='numbers')) # yapf: disable | |
| options = StandardOptions(streaming=True) | |
| p = TestPipeline(options=options) | |
| main = p | test_stream | |
| # Use an AfterWatermark trigger with an early firing to test that the | |
| # watermark is advancing properly and that the element is being emitted in | |
| # the correct window. | |
| letters = ( | |
| main['letters'] | |
| | 'letter windows' >> beam.WindowInto( | |
| FixedWindows(15), | |
| trigger=trigger.AfterWatermark(early=trigger.AfterCount(1)), | |
| accumulation_mode=trigger.AccumulationMode.DISCARDING) | |
| | 'letter with key' >> beam.Map(lambda x: ('k', x)) | |
| | 'letter gbk' >> beam.GroupByKey()) | |
| numbers = ( | |
| main['numbers'] | |
| | 'number windows' >> beam.WindowInto( | |
| FixedWindows(15), | |
| trigger=trigger.AfterWatermark(early=trigger.AfterCount(1)), | |
| accumulation_mode=trigger.AccumulationMode.DISCARDING) | |
| | 'number with key' >> beam.Map(lambda x: ('k', x)) | |
| | 'number gbk' >> beam.GroupByKey()) | |
| # The letters were emitted when the watermark was at 5, thus we expect to | |
| # see the elements in the [0, 15) window. We used an early trigger to make | |
| # sure that the ON_TIME empty pane was also emitted with a TestStream. | |
| # This pane has no data because of the early trigger causes the elements to | |
| # fire before the end of the window and because the accumulation mode | |
| # discards any data after the trigger fired. | |
| expected_letters = { | |
| window.IntervalWindow(0, 15): [ | |
| ('k', ['a', 'b', 'c']), | |
| ('k', []), | |
| ], | |
| } | |
| # Same here, except the numbers were emitted at watermark = 20, thus they | |
| # are in the [15, 30) window. | |
| expected_numbers = { | |
| window.IntervalWindow(15, 30): [ | |
| ('k', ['1', '2', '3']), | |
| ('k', []), | |
| ], | |
| } | |
| assert_that( | |
| letters, | |
| equal_to_per_window(expected_letters), | |
| label='letters assert per window') | |
| assert_that( | |
| numbers, | |
| equal_to_per_window(expected_numbers), | |
| label='numbers assert per window') | |
| p.run() | |
| def test_dicts_not_interpreted_as_windowed_values(self): | |
| test_stream = (TestStream() | |
| .advance_processing_time(10) | |
| .advance_watermark_to(10) | |
| .add_elements([{'a': 0, 'b': 1, 'c': 2}]) | |
| .advance_watermark_to_infinity()) # yapf: disable | |
| class RecordFn(beam.DoFn): | |
| def process( | |
| self, | |
| element=beam.DoFn.ElementParam, | |
| timestamp=beam.DoFn.TimestampParam): | |
| yield (element, timestamp) | |
| options = PipelineOptions() | |
| options.view_as(StandardOptions).streaming = True | |
| with TestPipeline(options=options) as p: | |
| my_record_fn = RecordFn() | |
| records = p | test_stream | beam.ParDo(my_record_fn) | |
| assert_that( | |
| records, | |
| equal_to([ | |
| ({ | |
| 'a': 0, 'b': 1, 'c': 2 | |
| }, timestamp.Timestamp(10)), | |
| ])) | |
| def test_windowed_values_interpreted_correctly(self): | |
| windowed_value = WindowedValueHolder( | |
| WindowedValue( | |
| 'a', | |
| Timestamp(5), [beam.window.IntervalWindow(5, 10)], | |
| PaneInfo(True, True, PaneInfoTiming.ON_TIME, 0, 0))) | |
| test_stream = (TestStream() | |
| .advance_processing_time(10) | |
| .advance_watermark_to(10) | |
| .add_elements([windowed_value]) | |
| .advance_watermark_to_infinity()) # yapf: disable | |
| class RecordFn(beam.DoFn): | |
| def process( | |
| self, | |
| element=beam.DoFn.ElementParam, | |
| timestamp=beam.DoFn.TimestampParam, | |
| window=beam.DoFn.WindowParam): | |
| yield (element, timestamp, window) | |
| options = PipelineOptions() | |
| options.view_as(StandardOptions).streaming = True | |
| with TestPipeline(options=options) as p: | |
| my_record_fn = RecordFn() | |
| records = p | test_stream | beam.ParDo(my_record_fn) | |
| assert_that( | |
| records, | |
| equal_to([ | |
| ('a', timestamp.Timestamp(5), beam.window.IntervalWindow(5, 10)), | |
| ])) | |
| def test_instance_check_windowed_value_holder(self): | |
| windowed_value = WindowedValue( | |
| 'a', | |
| Timestamp(5), [beam.window.IntervalWindow(5, 10)], | |
| PaneInfo(True, True, PaneInfoTiming.ON_TIME, 0, 0)) | |
| self.assertTrue( | |
| isinstance(WindowedValueHolder(windowed_value), WindowedValueHolder)) | |
| self.assertTrue( | |
| isinstance( | |
| beam.Row( | |
| windowed_value=windowed_value, urn=common_urns.coders.ROW.urn), | |
| WindowedValueHolder)) | |
| self.assertFalse( | |
| isinstance( | |
| beam.Row(windowed_value=windowed_value), WindowedValueHolder)) | |
| self.assertFalse(isinstance(windowed_value, WindowedValueHolder)) | |
| self.assertFalse( | |
| isinstance(beam.Row(x=windowed_value), WindowedValueHolder)) | |
| self.assertFalse( | |
| isinstance(beam.Row(windowed_value=1), WindowedValueHolder)) | |
| def test_gbk_execution_no_triggers(self): | |
| test_stream = ( | |
| TestStream().advance_watermark_to(10).add_elements([ | |
| 'a', 'b', 'c' | |
| ]).advance_watermark_to(20).add_elements(['d']).add_elements([ | |
| 'e' | |
| ]).advance_processing_time(10).advance_watermark_to(300).add_elements([ | |
| TimestampedValue('late', 12) | |
| ]).add_elements([TimestampedValue('last', 310) | |
| ]).advance_watermark_to_infinity()) | |
| options = PipelineOptions() | |
| options.view_as(StandardOptions).streaming = True | |
| p = TestPipeline(options=options) | |
| records = ( | |
| p | |
| | test_stream | |
| | beam.WindowInto(FixedWindows(15), allowed_lateness=300) | |
| | beam.Map(lambda x: ('k', x)) | |
| | beam.GroupByKey()) | |
| # TODO(https://github.com/apache/beam/issues/18441): timestamp assignment | |
| # for elements from a GBK should respect the TimestampCombiner. The test | |
| # below should also verify the timestamps of the outputted elements once | |
| # this is implemented. | |
| # assert per window | |
| expected_window_to_elements = { | |
| window.IntervalWindow(0, 15): [ | |
| ('k', ['a', 'b', 'c']), | |
| ('k', ['late']), | |
| ], | |
| window.IntervalWindow(15, 30): [ | |
| ('k', ['d', 'e']), | |
| ], | |
| window.IntervalWindow(300, 315): [ | |
| ('k', ['last']), | |
| ], | |
| } | |
| assert_that( | |
| records, | |
| equal_to_per_window(expected_window_to_elements), | |
| label='assert per window') | |
| p.run() | |
| def test_gbk_execution_after_watermark_trigger(self): | |
| test_stream = (TestStream() | |
| .advance_watermark_to(10) | |
| .add_elements([TimestampedValue('a', 11)]) | |
| .advance_watermark_to(20) | |
| .add_elements([TimestampedValue('b', 21)]) | |
| .advance_watermark_to_infinity()) # yapf: disable | |
| options = PipelineOptions() | |
| options.view_as(StandardOptions).streaming = True | |
| p = TestPipeline(options=options) | |
| records = ( | |
| p # pylint: disable=unused-variable | |
| | test_stream | |
| | beam.WindowInto( | |
| FixedWindows(15), | |
| trigger=trigger.AfterWatermark(early=trigger.AfterCount(1)), | |
| accumulation_mode=trigger.AccumulationMode.DISCARDING) | |
| | beam.Map(lambda x: ('k', x)) | |
| | beam.GroupByKey()) | |
| # TODO(https://github.com/apache/beam/issues/18441): timestamp assignment | |
| # for elements from a GBK should respect the TimestampCombiner. The test | |
| # below should also verify the timestamps of the outputted elements once | |
| # this is implemented. | |
| # assert per window | |
| expected_window_to_elements = { | |
| window.IntervalWindow(0, 15): [('k', ['a']), ('k', [])], | |
| window.IntervalWindow(15, 30): [('k', ['b']), ('k', [])], | |
| } | |
| assert_that( | |
| records, | |
| equal_to_per_window(expected_window_to_elements), | |
| label='assert per window') | |
| p.run() | |
| def test_gbk_execution_after_processing_trigger_fired(self): | |
| """Advance TestClock to (X + delta) and see the pipeline does finish.""" | |
| # TODO(mariagh): Add test_gbk_execution_after_processing_trigger_unfired | |
| # Advance TestClock to (X + delta) and see the pipeline does finish | |
| # Possibly to the framework trigger_transcripts.yaml | |
| test_stream = (TestStream() | |
| .advance_watermark_to(10) | |
| .add_elements(['a']) | |
| .advance_processing_time(5.1) | |
| .advance_watermark_to_infinity()) # yapf: disable | |
| options = PipelineOptions() | |
| options.view_as(StandardOptions).streaming = True | |
| options.view_as(TypeOptions).allow_unsafe_triggers = True | |
| p = TestPipeline(options=options) | |
| records = ( | |
| p | |
| | test_stream | |
| | beam.WindowInto( | |
| beam.window.FixedWindows(15), | |
| trigger=trigger.AfterProcessingTime(5), | |
| accumulation_mode=trigger.AccumulationMode.DISCARDING) | |
| | beam.Map(lambda x: ('k', x)) | |
| | beam.GroupByKey()) | |
| # TODO(https://github.com/apache/beam/issues/18441): timestamp assignment | |
| # for elements from a GBK should respect the TimestampCombiner. The test | |
| # below should also verify the timestamps of the outputted elements once | |
| # this is implemented. | |
| expected_window_to_elements = { | |
| window.IntervalWindow(0, 15): [('k', ['a'])], | |
| } | |
| assert_that( | |
| records, | |
| equal_to_per_window(expected_window_to_elements), | |
| label='assert per window') | |
| p.run() | |
| def test_basic_execution_batch_sideinputs(self): | |
| options = PipelineOptions() | |
| options.view_as(StandardOptions).streaming = True | |
| p = TestPipeline(options=options) | |
| main_stream = (p | |
| | 'main TestStream' >> TestStream() | |
| .advance_watermark_to(10) | |
| .add_elements(['e']) | |
| .advance_watermark_to_infinity()) # yapf: disable | |
| side = ( | |
| p | |
| | beam.Create([2, 1, 4]) | |
| | beam.Map(lambda t: window.TimestampedValue(t, t))) | |
| class RecordFn(beam.DoFn): | |
| def process( | |
| self, | |
| elm=beam.DoFn.ElementParam, | |
| ts=beam.DoFn.TimestampParam, | |
| side=beam.DoFn.SideInputParam): | |
| yield (elm, ts, sorted(side)) | |
| records = ( | |
| main_stream # pylint: disable=unused-variable | |
| | beam.ParDo(RecordFn(), beam.pvalue.AsList(side))) | |
| assert_that(records, equal_to([('e', Timestamp(10), [1, 2, 4])])) | |
| p.run() | |
| def test_basic_execution_sideinputs(self): | |
| options = PipelineOptions() | |
| options.view_as(StandardOptions).streaming = True | |
| with TestPipeline(options=options) as p: | |
| test_stream = (p | TestStream() | |
| .advance_watermark_to(0, tag='side') | |
| .advance_watermark_to(10, tag='main') | |
| .add_elements(['e'], tag='main') | |
| .add_elements([window.TimestampedValue(2, 2)], tag='side') | |
| .add_elements([window.TimestampedValue(1, 1)], tag='side') | |
| .add_elements([window.TimestampedValue(7, 7)], tag='side') | |
| .add_elements([window.TimestampedValue(4, 4)], tag='side') | |
| ) # yapf: disable | |
| main_stream = test_stream['main'] | |
| side_stream = test_stream['side'] | |
| class RecordFn(beam.DoFn): | |
| def process( | |
| self, | |
| elm=beam.DoFn.ElementParam, | |
| ts=beam.DoFn.TimestampParam, | |
| side=beam.DoFn.SideInputParam): | |
| yield (elm, ts, side) | |
| records = ( | |
| main_stream # pylint: disable=unused-variable | |
| | beam.ParDo(RecordFn(), beam.pvalue.AsList(side_stream))) | |
| assert_that(records, equal_to([('e', Timestamp(10), [2, 1, 7, 4])])) | |
| def test_basic_execution_batch_sideinputs_fixed_windows(self): | |
| options = PipelineOptions() | |
| options.view_as(StandardOptions).streaming = True | |
| p = TestPipeline(options=options) | |
| main_stream = ( | |
| p | |
| | | |
| 'main TestStream' >> TestStream().advance_watermark_to(2).add_elements( | |
| ['a']).advance_watermark_to(4).add_elements( | |
| ['b']).advance_watermark_to_infinity() | |
| | 'main window' >> beam.WindowInto(window.FixedWindows(1))) | |
| side = ( | |
| p | |
| | beam.Create([2, 1, 4]) | |
| | beam.Map(lambda t: window.TimestampedValue(t, t)) | |
| | beam.WindowInto(window.FixedWindows(2))) | |
| class RecordFn(beam.DoFn): | |
| def process( | |
| self, | |
| elm=beam.DoFn.ElementParam, | |
| ts=beam.DoFn.TimestampParam, | |
| side=beam.DoFn.SideInputParam): | |
| yield (elm, ts, side) | |
| records = ( | |
| main_stream # pylint: disable=unused-variable | |
| | beam.ParDo(RecordFn(), beam.pvalue.AsList(side))) | |
| # assert per window | |
| expected_window_to_elements = { | |
| window.IntervalWindow(2, 3): [('a', Timestamp(2), [2])], | |
| window.IntervalWindow(4, 5): [('b', Timestamp(4), [4])] | |
| } | |
| assert_that( | |
| records, | |
| equal_to_per_window(expected_window_to_elements), | |
| label='assert per window') | |
| p.run() | |
| def test_basic_execution_sideinputs_fixed_windows(self): | |
| options = PipelineOptions() | |
| options.view_as(StandardOptions).streaming = True | |
| p = TestPipeline(options=options) | |
| test_stream = (p | TestStream() | |
| .advance_watermark_to(12, tag='side') | |
| .add_elements([window.TimestampedValue('s1', 10)], tag='side') | |
| .advance_watermark_to(20, tag='side') | |
| .add_elements([window.TimestampedValue('s2', 20)], tag='side') | |
| .advance_watermark_to(9, tag='main') | |
| .add_elements(['a1', 'a2', 'a3', 'a4'], tag='main') | |
| .add_elements(['b'], tag='main') | |
| .advance_watermark_to(18, tag='main') | |
| .add_elements('c', tag='main') | |
| ) # yapf: disable | |
| main_stream = ( | |
| test_stream['main'] | |
| | 'main windowInto' >> beam.WindowInto(window.FixedWindows(1))) | |
| side_stream = ( | |
| test_stream['side'] | |
| | 'side windowInto' >> beam.WindowInto(window.FixedWindows(3))) | |
| class RecordFn(beam.DoFn): | |
| def process( | |
| self, | |
| elm=beam.DoFn.ElementParam, | |
| ts=beam.DoFn.TimestampParam, | |
| side=beam.DoFn.SideInputParam): | |
| yield (elm, ts, side) | |
| records = ( | |
| main_stream # pylint: disable=unused-variable | |
| | beam.ParDo(RecordFn(), beam.pvalue.AsList(side_stream))) | |
| # assert per window | |
| expected_window_to_elements = { | |
| window.IntervalWindow(9, 10): [ | |
| ('a1', Timestamp(9), ['s1']), ('a2', Timestamp(9), ['s1']), | |
| ('a3', Timestamp(9), ['s1']), ('a4', Timestamp(9), ['s1']), | |
| ('b', Timestamp(9), ['s1']) | |
| ], | |
| window.IntervalWindow(18, 19): [('c', Timestamp(18), ['s2'])], | |
| } | |
| assert_that( | |
| records, | |
| equal_to_per_window(expected_window_to_elements), | |
| label='assert per window') | |
| p.run() | |
| def test_roundtrip_proto(self): | |
| test_stream = (TestStream() | |
| .advance_processing_time(1) | |
| .advance_watermark_to(2) | |
| .add_elements([1, 2, 3])) # yapf: disable | |
| p = TestPipeline(options=StandardOptions(streaming=True)) | |
| p | test_stream | |
| pipeline_proto, context = p.to_runner_api(return_context=True) | |
| for t in pipeline_proto.components.transforms.values(): | |
| if t.spec.urn == common_urns.primitives.TEST_STREAM.urn: | |
| test_stream_proto = t | |
| self.assertTrue(test_stream_proto) | |
| roundtrip_test_stream = TestStream().from_runner_api( | |
| test_stream_proto, context) | |
| self.assertListEqual(test_stream._events, roundtrip_test_stream._events) | |
| self.assertSetEqual( | |
| test_stream.output_tags, roundtrip_test_stream.output_tags) | |
| self.assertEqual(test_stream.coder, roundtrip_test_stream.coder) | |
| def test_roundtrip_proto_multi(self): | |
| test_stream = (TestStream() | |
| .advance_processing_time(1) | |
| .advance_watermark_to(2, tag='a') | |
| .advance_watermark_to(3, tag='b') | |
| .add_elements([1, 2, 3], tag='a') | |
| .add_elements([4, 5, 6], tag='b')) # yapf: disable | |
| options = StandardOptions(streaming=True) | |
| p = TestPipeline(options=options) | |
| p | test_stream | |
| pipeline_proto, context = p.to_runner_api(return_context=True) | |
| for t in pipeline_proto.components.transforms.values(): | |
| if t.spec.urn == common_urns.primitives.TEST_STREAM.urn: | |
| test_stream_proto = t | |
| self.assertTrue(test_stream_proto) | |
| roundtrip_test_stream = TestStream().from_runner_api( | |
| test_stream_proto, context) | |
| self.assertListEqual(test_stream._events, roundtrip_test_stream._events) | |
| self.assertSetEqual( | |
| test_stream.output_tags, roundtrip_test_stream.output_tags) | |
| self.assertEqual(test_stream.coder, roundtrip_test_stream.coder) | |
| def test_basic_execution_with_service(self): | |
| """Tests that the TestStream can correctly read from an RPC service. | |
| """ | |
| coder = beam.coders.FastPrimitivesCoder() | |
| test_stream_events = (TestStream(coder=coder) | |
| .advance_watermark_to(10000) | |
| .add_elements(['a', 'b', 'c']) | |
| .advance_watermark_to(20000) | |
| .add_elements(['d']) | |
| .add_elements(['e']) | |
| .advance_processing_time(10) | |
| .advance_watermark_to(300000) | |
| .add_elements([TimestampedValue('late', 12000)]) | |
| .add_elements([TimestampedValue('last', 310000)]) | |
| .advance_watermark_to_infinity())._events # yapf: disable | |
| test_stream_proto_events = [ | |
| e.to_runner_api(coder) for e in test_stream_events | |
| ] | |
| class InMemoryEventReader: | |
| def read_multiple(self, unused_keys): | |
| for e in test_stream_proto_events: | |
| yield e | |
| service = TestStreamServiceController(reader=InMemoryEventReader()) | |
| service.start() | |
| test_stream = TestStream(coder=coder, endpoint=service.endpoint) | |
| class RecordFn(beam.DoFn): | |
| def process( | |
| self, | |
| element=beam.DoFn.ElementParam, | |
| timestamp=beam.DoFn.TimestampParam): | |
| yield (element, timestamp) | |
| options = StandardOptions(streaming=True) | |
| p = TestPipeline(options=options) | |
| my_record_fn = RecordFn() | |
| records = p | test_stream | beam.ParDo(my_record_fn) | |
| assert_that( | |
| records, | |
| equal_to([ | |
| ('a', timestamp.Timestamp(10)), | |
| ('b', timestamp.Timestamp(10)), | |
| ('c', timestamp.Timestamp(10)), | |
| ('d', timestamp.Timestamp(20)), | |
| ('e', timestamp.Timestamp(20)), | |
| ('late', timestamp.Timestamp(12)), | |
| ('last', timestamp.Timestamp(310)), | |
| ])) | |
| p.run() | |
| class ReverseTestStreamTest(unittest.TestCase): | |
| def test_basic_execution(self): | |
| test_stream = (TestStream() | |
| .advance_watermark_to(0) | |
| .advance_processing_time(5) | |
| .add_elements(['a', 'b', 'c']) | |
| .advance_watermark_to(2) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(4) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(6) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(8) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(10) | |
| .advance_processing_time(1) | |
| .add_elements([TimestampedValue('1', 15), | |
| TimestampedValue('2', 15), | |
| TimestampedValue('3', 15)])) # yapf: disable | |
| options = StandardOptions(streaming=True) | |
| p = TestPipeline(options=options) | |
| records = ( | |
| p | |
| | test_stream | |
| | ReverseTestStream(sample_resolution_sec=1, output_tag=None)) | |
| assert_that( | |
| records, | |
| equal_to_per_window({ | |
| beam.window.GlobalWindow(): [ | |
| [ProcessingTimeEvent(5), WatermarkEvent(0)], | |
| [ | |
| ElementEvent([ | |
| TimestampedValue('a', 0), | |
| TimestampedValue('b', 0), | |
| TimestampedValue('c', 0) | |
| ]) | |
| ], | |
| [ProcessingTimeEvent(1), WatermarkEvent(2000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(4000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(6000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(8000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(10000000)], | |
| [ | |
| ElementEvent([ | |
| TimestampedValue('1', 15), | |
| TimestampedValue('2', 15), | |
| TimestampedValue('3', 15) | |
| ]) | |
| ], | |
| ], | |
| })) | |
| p.run() | |
| def test_windowing(self): | |
| test_stream = (TestStream() | |
| .advance_watermark_to(0) | |
| .add_elements(['a', 'b', 'c']) | |
| .advance_processing_time(1) | |
| .advance_processing_time(1) | |
| .advance_processing_time(1) | |
| .advance_processing_time(1) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(5) | |
| .add_elements(['1', '2', '3']) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(6) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(7) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(8) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(9) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(10) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(11) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(12) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(13) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(14) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(15) | |
| .advance_processing_time(1) | |
| ) # yapf: disable | |
| options = StandardOptions(streaming=True) | |
| p = TestPipeline(options=options) | |
| records = ( | |
| p | |
| | test_stream | |
| | 'letter windows' >> beam.WindowInto( | |
| FixedWindows(5), | |
| accumulation_mode=trigger.AccumulationMode.DISCARDING) | |
| | 'letter with key' >> beam.Map(lambda x: ('k', x)) | |
| | 'letter gbk' >> beam.GroupByKey() | |
| | ReverseTestStream(sample_resolution_sec=1, output_tag=None)) | |
| assert_that( | |
| records, | |
| equal_to_per_window({ | |
| beam.window.GlobalWindow(): [ | |
| [ProcessingTimeEvent(5), WatermarkEvent(4999998)], | |
| [ | |
| ElementEvent( | |
| [TimestampedValue(('k', ['a', 'b', 'c']), 4.999999)]) | |
| ], | |
| [ProcessingTimeEvent(1), WatermarkEvent(5000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(6000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(7000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(8000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(9000000)], | |
| [ | |
| ElementEvent( | |
| [TimestampedValue(('k', ['1', '2', '3']), 9.999999)]) | |
| ], | |
| [ProcessingTimeEvent(1), WatermarkEvent(10000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(11000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(12000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(13000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(14000000)], | |
| [ProcessingTimeEvent(1), WatermarkEvent(15000000)], | |
| ], | |
| })) | |
| p.run() | |
| def test_basic_execution_in_records_format(self): | |
| test_stream = (TestStream() | |
| .advance_watermark_to(0) | |
| .advance_processing_time(5) | |
| .add_elements(['a', 'b', 'c']) | |
| .advance_watermark_to(2) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(4) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(6) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(8) | |
| .advance_processing_time(1) | |
| .advance_watermark_to(10) | |
| .advance_processing_time(1) | |
| .add_elements([TimestampedValue('1', 15), | |
| TimestampedValue('2', 15), | |
| TimestampedValue('3', 15)])) # yapf: disable | |
| options = StandardOptions(streaming=True) | |
| p = TestPipeline(options=options) | |
| coder = beam.coders.FastPrimitivesCoder() | |
| records = ( | |
| p | |
| | test_stream | |
| | ReverseTestStream( | |
| sample_resolution_sec=1, | |
| coder=coder, | |
| output_format=OutputFormat.TEST_STREAM_FILE_RECORDS, | |
| output_tag=None) | |
| | 'stringify' >> beam.Map(str)) | |
| assert_that( | |
| records, | |
| equal_to_per_window({ | |
| beam.window.GlobalWindow(): [ | |
| str(beam_interactive_api_pb2.TestStreamFileHeader()), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| processing_time_event=beam_runner_api_pb2. | |
| TestStreamPayload.Event.AdvanceProcessingTime( | |
| advance_duration=5000000)))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| watermark_event=beam_runner_api_pb2. | |
| TestStreamPayload.Event.AdvanceWatermark( | |
| new_watermark=0)))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| element_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event.AddElements( | |
| elements=[ | |
| beam_runner_api_pb2.TestStreamPayload. | |
| TimestampedElement( | |
| encoded_element=coder.encode('a'), | |
| timestamp=0), | |
| beam_runner_api_pb2.TestStreamPayload. | |
| TimestampedElement( | |
| encoded_element=coder.encode('b'), | |
| timestamp=0), | |
| beam_runner_api_pb2.TestStreamPayload. | |
| TimestampedElement( | |
| encoded_element=coder.encode('c'), | |
| timestamp=0), | |
| ])))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| watermark_event=beam_runner_api_pb2. | |
| TestStreamPayload.Event.AdvanceWatermark( | |
| new_watermark=2000000)))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| processing_time_event=beam_runner_api_pb2. | |
| TestStreamPayload.Event.AdvanceProcessingTime( | |
| advance_duration=1000000)))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| watermark_event=beam_runner_api_pb2. | |
| TestStreamPayload.Event.AdvanceWatermark( | |
| new_watermark=4000000)))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| processing_time_event=beam_runner_api_pb2. | |
| TestStreamPayload.Event.AdvanceProcessingTime( | |
| advance_duration=1000000)))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| watermark_event=beam_runner_api_pb2. | |
| TestStreamPayload.Event.AdvanceWatermark( | |
| new_watermark=6000000)))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| processing_time_event=beam_runner_api_pb2. | |
| TestStreamPayload.Event.AdvanceProcessingTime( | |
| advance_duration=1000000)))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| watermark_event=beam_runner_api_pb2. | |
| TestStreamPayload.Event.AdvanceWatermark( | |
| new_watermark=8000000)))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| processing_time_event=beam_runner_api_pb2. | |
| TestStreamPayload.Event.AdvanceProcessingTime( | |
| advance_duration=1000000)))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| watermark_event=beam_runner_api_pb2. | |
| TestStreamPayload.Event.AdvanceWatermark( | |
| new_watermark=10000000)))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| processing_time_event=beam_runner_api_pb2. | |
| TestStreamPayload.Event.AdvanceProcessingTime( | |
| advance_duration=1000000)))), | |
| str( | |
| beam_interactive_api_pb2.TestStreamFileRecord( | |
| recorded_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event( | |
| element_event=beam_runner_api_pb2.TestStreamPayload. | |
| Event.AddElements( | |
| elements=[ | |
| beam_runner_api_pb2.TestStreamPayload. | |
| TimestampedElement( | |
| encoded_element=coder.encode('1'), | |
| timestamp=15000000), | |
| beam_runner_api_pb2.TestStreamPayload. | |
| TimestampedElement( | |
| encoded_element=coder.encode('2'), | |
| timestamp=15000000), | |
| beam_runner_api_pb2.TestStreamPayload. | |
| TimestampedElement( | |
| encoded_element=coder.encode('3'), | |
| timestamp=15000000), | |
| ])))), | |
| ], | |
| })) | |
| p.run() | |
| if __name__ == '__main__': | |
| unittest.main() |