github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/io/watermark_estimators_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 built-in WatermarkEstimators"""
    19  
    20  # pytype: skip-file
    21  
    22  import unittest
    23  
    24  import mock
    25  
    26  from apache_beam.io.iobase import WatermarkEstimator
    27  from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
    28  from apache_beam.io.watermark_estimators import MonotonicWatermarkEstimator
    29  from apache_beam.io.watermark_estimators import WalltimeWatermarkEstimator
    30  from apache_beam.utils.timestamp import Duration
    31  from apache_beam.utils.timestamp import Timestamp
    32  
    33  
    34  class MonotonicWatermarkEstimatorTest(unittest.TestCase):
    35    def test_initialize_from_state(self):
    36      timestamp = Timestamp(10)
    37      watermark_estimator = MonotonicWatermarkEstimator(timestamp)
    38      self.assertIsInstance(watermark_estimator, WatermarkEstimator)
    39      self.assertEqual(watermark_estimator.get_estimator_state(), timestamp)
    40  
    41    def test_observe_timestamp(self):
    42      watermark_estimator = MonotonicWatermarkEstimator(Timestamp(10))
    43      watermark_estimator.observe_timestamp(Timestamp(15))
    44      self.assertEqual(watermark_estimator.current_watermark(), Timestamp(15))
    45      watermark_estimator.observe_timestamp(Timestamp(20))
    46      self.assertEqual(watermark_estimator.current_watermark(), Timestamp(20))
    47      watermark_estimator.observe_timestamp(Timestamp(20))
    48      self.assertEqual(watermark_estimator.current_watermark(), Timestamp(20))
    49      watermark_estimator.observe_timestamp(Timestamp(10))
    50      self.assertEqual(watermark_estimator.current_watermark(), Timestamp(20))
    51  
    52    def test_get_estimator_state(self):
    53      watermark_estimator = MonotonicWatermarkEstimator(Timestamp(10))
    54      self.assertEqual(watermark_estimator.get_estimator_state(), Timestamp(10))
    55      watermark_estimator.observe_timestamp(Timestamp(15))
    56      # State only progresses when we request a new watermark
    57      self.assertEqual(watermark_estimator.get_estimator_state(), Timestamp(10))
    58      self.assertEqual(watermark_estimator.current_watermark(), Timestamp(15))
    59      self.assertEqual(watermark_estimator.get_estimator_state(), Timestamp(15))
    60  
    61  
    62  class WalltimeWatermarkEstimatorTest(unittest.TestCase):
    63    @mock.patch('apache_beam.utils.timestamp.Timestamp.now')
    64    def test_initialization(self, mock_timestamp):
    65      now_time = Timestamp.now() - Duration(10)
    66      mock_timestamp.side_effect = lambda: now_time
    67      watermark_estimator = WalltimeWatermarkEstimator()
    68      self.assertIsInstance(watermark_estimator, WatermarkEstimator)
    69      self.assertEqual(watermark_estimator.get_estimator_state(), now_time)
    70  
    71    def test_observe_timestamp(self):
    72      now_time = Timestamp.now() + Duration(10)
    73      watermark_estimator = WalltimeWatermarkEstimator(now_time)
    74      watermark_estimator.observe_timestamp(Timestamp(10))
    75      watermark_estimator.observe_timestamp(Timestamp(10))
    76      self.assertEqual(watermark_estimator.current_watermark(), now_time)
    77  
    78    def test_advance_watermark_with_incorrect_sys_clock(self):
    79      initial_timestamp = Timestamp.now() + Duration(100)
    80      watermark_estimator = WalltimeWatermarkEstimator(initial_timestamp)
    81      self.assertEqual(watermark_estimator.current_watermark(), initial_timestamp)
    82      self.assertEqual(
    83          watermark_estimator.get_estimator_state(), initial_timestamp)
    84  
    85  
    86  class ManualWatermarkEstimatorTest(unittest.TestCase):
    87    def test_initialization(self):
    88      watermark_estimator = ManualWatermarkEstimator(None)
    89      self.assertIsNone(watermark_estimator.get_estimator_state())
    90      self.assertIsNone(watermark_estimator.current_watermark())
    91      watermark_estimator = ManualWatermarkEstimator(Timestamp(10))
    92      self.assertEqual(watermark_estimator.get_estimator_state(), Timestamp(10))
    93  
    94    def test_set_watermark(self):
    95      watermark_estimator = ManualWatermarkEstimator(None)
    96      self.assertIsNone(watermark_estimator.current_watermark())
    97      watermark_estimator.observe_timestamp(Timestamp(10))
    98      self.assertIsNone(watermark_estimator.current_watermark())
    99      watermark_estimator.set_watermark(Timestamp(20))
   100      self.assertEqual(watermark_estimator.current_watermark(), Timestamp(20))
   101      watermark_estimator.set_watermark(Timestamp(30))
   102      self.assertEqual(watermark_estimator.current_watermark(), Timestamp(30))
   103      with self.assertRaises(ValueError):
   104        watermark_estimator.set_watermark(Timestamp(25))
   105  
   106  
   107  if __name__ == '__main__':
   108    unittest.main()