github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/transforms/periodicsequence_test.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """Unit tests for the PTransform and descendants."""
    19  
    20  # pytype: skip-file
    21  
    22  import inspect
    23  import time
    24  import unittest
    25  
    26  import apache_beam as beam
    27  from apache_beam.testing.test_pipeline import TestPipeline
    28  from apache_beam.testing.util import assert_that
    29  from apache_beam.testing.util import equal_to
    30  from apache_beam.transforms.periodicsequence import PeriodicImpulse
    31  from apache_beam.transforms.periodicsequence import PeriodicSequence
    32  
    33  # Disable frequent lint warning due to pipe operator for chaining transforms.
    34  # pylint: disable=expression-not-assigned
    35  
    36  
    37  class PeriodicSequenceTest(unittest.TestCase):
    38    def test_periodicsequence_outputs_valid_sequence(self):
    39      start_offset = 1
    40      start_time = time.time() + start_offset
    41      duration = 1
    42      end_time = start_time + duration
    43      interval = 0.25
    44  
    45      with TestPipeline() as p:
    46        result = (
    47            p
    48            | 'ImpulseElement' >> beam.Create([(start_time, end_time, interval)])
    49            | 'ImpulseSeqGen' >> PeriodicSequence())
    50  
    51        k = [
    52            start_time + x * interval
    53            for x in range(0, int(duration / interval), 1)
    54        ]
    55        self.assertEqual(result.is_bounded, False)
    56        assert_that(result, equal_to(k))
    57  
    58    def test_periodicimpulse_windowing_on_si(self):
    59      start_offset = -15
    60      it = time.time() + start_offset
    61      duration = 15
    62      et = it + duration
    63      interval = 5
    64  
    65      with TestPipeline() as p:
    66        si = (
    67            p
    68            | 'PeriodicImpulse' >> PeriodicImpulse(it, et, interval, True)
    69            | 'AddKey' >> beam.Map(lambda v: ('key', v))
    70            | 'GBK' >> beam.GroupByKey()
    71            | 'SortGBK' >> beam.MapTuple(lambda k, vs: (k, sorted(vs))))
    72  
    73        actual = si
    74        k = [('key', [it + x * interval])
    75             for x in range(0, int(duration / interval), 1)]
    76        assert_that(actual, equal_to(k))
    77  
    78    def test_periodicimpulse_default_start(self):
    79      default_parameters = inspect.signature(PeriodicImpulse.__init__).parameters
    80      it = default_parameters["start_timestamp"].default
    81      duration = 1
    82      et = it + duration
    83      interval = 0.5
    84  
    85      # Check default `stop_timestamp` is the same type `start_timestamp`
    86      is_same_type = isinstance(
    87          it, type(default_parameters["stop_timestamp"].default))
    88      error = "'start_timestamp' and 'stop_timestamp' have different type"
    89      assert is_same_type, error
    90  
    91      with TestPipeline() as p:
    92        result = p | 'PeriodicImpulse' >> PeriodicImpulse(it, et, interval)
    93  
    94        k = [it + x * interval for x in range(0, int(duration / interval))]
    95        self.assertEqual(result.is_bounded, False)
    96        assert_that(result, equal_to(k))
    97  
    98    def test_periodicsequence_outputs_valid_sequence_in_past(self):
    99      start_offset = -10000
   100      it = time.time() + start_offset
   101      duration = 5
   102      et = it + duration
   103      interval = 1
   104  
   105      with TestPipeline() as p:
   106        result = (
   107            p
   108            | 'ImpulseElement' >> beam.Create([(it, et, interval)])
   109            | 'ImpulseSeqGen' >> PeriodicSequence())
   110  
   111        k = [it + x * interval for x in range(0, int(duration / interval), 1)]
   112        self.assertEqual(result.is_bounded, False)
   113        assert_that(result, equal_to(k))
   114  
   115  
   116  if __name__ == '__main__':
   117    unittest.main()