github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/utils/counters_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 counters and counter names."""
    19  
    20  # pytype: skip-file
    21  
    22  import unittest
    23  
    24  from parameterized import parameterized_class
    25  
    26  from apache_beam.utils import counters
    27  from apache_beam.utils.counters import CounterName
    28  
    29  
    30  class CounterNameTest(unittest.TestCase):
    31    def test_name_string_representation(self):
    32      counter_name = CounterName('counter_name', 'stage_name', 'step_name')
    33  
    34      # This string representation is utilized by the worker to report progress.
    35      # Change only if the worker code has also been changed.
    36      self.assertEqual('stage_name-step_name-counter_name', str(counter_name))
    37      self.assertIn(
    38          '<CounterName<stage_name-step_name-counter_name> at 0x',
    39          repr(counter_name))
    40  
    41    def test_equal_objects(self):
    42      self.assertEqual(
    43          CounterName('counter_name', 'stage_name', 'step_name'),
    44          CounterName('counter_name', 'stage_name', 'step_name'))
    45      self.assertNotEqual(
    46          CounterName('counter_name', 'stage_name', 'step_name'),
    47          CounterName('counter_name', 'stage_name', 'step_nam'))
    48  
    49      # Testing objects with an IOTarget.
    50      self.assertEqual(
    51          CounterName(
    52              'counter_name',
    53              'stage_name',
    54              'step_name',
    55              io_target=counters.side_input_id(1, 's9')),
    56          CounterName(
    57              'counter_name',
    58              'stage_name',
    59              'step_name',
    60              io_target=counters.side_input_id(1, 's9')))
    61      self.assertNotEqual(
    62          CounterName(
    63              'counter_name',
    64              'stage_name',
    65              'step_name',
    66              io_target=counters.side_input_id(1, 's')),
    67          CounterName(
    68              'counter_name',
    69              'stage_name',
    70              'step_name',
    71              io_target=counters.side_input_id(1, 's9')))
    72  
    73    def test_hash_two_objects(self):
    74      self.assertEqual(
    75          hash(CounterName('counter_name', 'stage_name', 'step_name')),
    76          hash(CounterName('counter_name', 'stage_name', 'step_name')))
    77      self.assertNotEqual(
    78          hash(CounterName('counter_name', 'stage_name', 'step_name')),
    79          hash(CounterName('counter_name', 'stage_name', 'step_nam')))
    80  
    81  
    82  class CounterTest(unittest.TestCase):
    83    def setUp(self):
    84      self.counter_factory = counters.CounterFactory()
    85  
    86    def test_sum_counter(self):
    87      sum_counter = self.counter_factory.get_counter(
    88          CounterName('sum', 'stage_foo', 'step_bar'), counters.Counter.SUM)
    89      for i in range(100):
    90        sum_counter.update(i)
    91  
    92      self.assertEqual(99 * 50, sum_counter.value())
    93  
    94    def test_mean_counter(self):
    95      mean_counter = self.counter_factory.get_counter(
    96          CounterName('mean', 'stage_foo', 'step_bar'), counters.Counter.MEAN)
    97      for i in range(100):
    98        mean_counter.update(i)
    99  
   100      self.assertEqual(49, mean_counter.value())
   101  
   102    def test_distribution_counter(self):
   103      distribution_counter = self.counter_factory.get_counter(
   104          CounterName('distribution', 'stage_foo', 'step_bar'),
   105          counters.Counter.BEAM_DISTRIBUTION)
   106      for i in range(100):
   107        distribution_counter.update(i)
   108  
   109      self.assertEqual((49, 4950, 100, 0, 99), distribution_counter.value())
   110  
   111  
   112  @parameterized_class([
   113      {
   114          'combiner': counters.Counter.SUM
   115      },
   116      {
   117          'combiner': counters.Counter.MEAN
   118      },
   119      {
   120          'combiner': counters.Counter.BEAM_DISTRIBUTION
   121      },
   122      {
   123          'combiner': counters.Counter.DATAFLOW_DISTRIBUTION
   124      },
   125  ])
   126  class GeneralCounterTest(unittest.TestCase):
   127    def setUp(self):
   128      self.counter_factory = counters.CounterFactory()
   129  
   130    def test_reset(self):
   131      counter = self.counter_factory.get_counter(
   132          CounterName(self.combiner.default_label, 'stage_foo', 'reset'),
   133          self.combiner)
   134  
   135      for value in range(100):
   136        counter.update(value)
   137      expected = counter.value()
   138      counter.reset()
   139  
   140      for value in range(100):
   141        counter.update(value)
   142  
   143      self.assertEqual(expected, counter.value())
   144  
   145    def test_update_n(self):
   146      counter = self.counter_factory.get_counter(
   147          CounterName(self.combiner.default_label, 'stage_foo', 'update_n'),
   148          self.combiner)
   149      for i in range(100):
   150        value = i
   151        n = 100 - i
   152        for _ in range(n):
   153          counter.update(value)
   154  
   155      expected = counter.value()
   156  
   157      counter.reset()
   158  
   159      for i in range(100):
   160        value = i
   161        n = 100 - i
   162        counter.update_n(value, n)
   163  
   164      self.assertEqual(expected, counter.value())
   165  
   166  
   167  if __name__ == '__main__':
   168    unittest.main()