github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/runners/interactive/options/capture_limiters_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  import unittest
    19  
    20  import pandas as pd
    21  
    22  from apache_beam.portability.api import beam_runner_api_pb2
    23  from apache_beam.runners.interactive.options.capture_limiters import CountLimiter
    24  from apache_beam.runners.interactive.options.capture_limiters import ProcessingTimeLimiter
    25  from apache_beam.utils.windowed_value import WindowedValue
    26  
    27  
    28  class CaptureLimitersTest(unittest.TestCase):
    29    def test_count_limiter(self):
    30      limiter = CountLimiter(5)
    31  
    32      for e in range(4):
    33        limiter.update(e)
    34  
    35      self.assertFalse(limiter.is_triggered())
    36      limiter.update(4)
    37      self.assertTrue(limiter.is_triggered())
    38  
    39    def test_count_limiter_with_dataframes(self):
    40      limiter = CountLimiter(5)
    41  
    42      # Test that empty dataframes don't count.
    43      for _ in range(10):
    44        df = WindowedValue(pd.DataFrame(), 0, [])
    45        limiter.update(df)
    46  
    47      self.assertFalse(limiter.is_triggered())
    48      df = WindowedValue(pd.DataFrame({'col': list(range(10))}), 0, [])
    49      limiter.update(df)
    50      self.assertTrue(limiter.is_triggered())
    51  
    52    def test_processing_time_limiter(self):
    53      limiter = ProcessingTimeLimiter(max_duration_secs=2)
    54  
    55      e = beam_runner_api_pb2.TestStreamPayload.Event()
    56      e.processing_time_event.advance_duration = int(1 * 1e6)
    57      limiter.update(e)
    58      self.assertFalse(limiter.is_triggered())
    59  
    60      e = beam_runner_api_pb2.TestStreamPayload.Event()
    61      e.processing_time_event.advance_duration = int(2 * 1e6)
    62      limiter.update(e)
    63      self.assertTrue(limiter.is_triggered())
    64  
    65  
    66  if __name__ == '__main__':
    67    unittest.main()