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

     1  # the License.  You may obtain a copy of the License at
     2  #
     3  #    http://www.apache.org/licenses/LICENSE-2.0
     4  #
     5  # Unless required by applicable law or agreed to in writing, software
     6  # distributed under the License is distributed on an "AS IS" BASIS,
     7  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     8  # See the License for the specific language governing permissions and
     9  # limitations under the License.
    10  #
    11  
    12  """Unit tests for DataflowDistributionCounter
    13  When Cython is available, unit tests will test on cythonized module,
    14  otherwise, test on pure python module
    15  """
    16  
    17  # pytype: skip-file
    18  
    19  import unittest
    20  
    21  from mock import Mock
    22  
    23  from apache_beam.transforms import DataflowDistributionCounter
    24  
    25  INT64_MAX = 2**63 - 1
    26  
    27  
    28  class DataflowDistributionAccumulatorTest(unittest.TestCase):
    29    def test_calculate_bucket_index_with_input_0(self):
    30      counter = DataflowDistributionCounter()
    31      index = counter.calculate_bucket_index(0)
    32      self.assertEqual(index, 0)
    33  
    34    def test_calculate_bucket_index_within_max_long(self):
    35      counter = DataflowDistributionCounter()
    36      bucket = 1
    37      power_of_ten = 1
    38      while power_of_ten <= INT64_MAX:
    39        for multiplier in [1, 2, 5]:
    40          value = multiplier * power_of_ten
    41          actual_bucket = counter.calculate_bucket_index(value - 1)
    42          self.assertEqual(actual_bucket, bucket - 1)
    43          bucket += 1
    44        power_of_ten *= 10
    45  
    46    def test_add_input(self):
    47      counter = DataflowDistributionCounter()
    48      expected_buckets = [1, 3, 0, 0, 0, 0, 0, 0, 1, 1]
    49      expected_sum = 1510
    50      expected_first_bucket_index = 1
    51      expected_count = 6
    52      expected_min = 1
    53      expected_max = 1000
    54      for element in [1, 500, 2, 3, 1000, 4]:
    55        counter.add_input(element)
    56      histogram = Mock(firstBucketOffset=None, bucketCounts=None)
    57      counter.translate_to_histogram(histogram)
    58      self.assertEqual(counter.sum, expected_sum)
    59      self.assertEqual(counter.count, expected_count)
    60      self.assertEqual(counter.min, expected_min)
    61      self.assertEqual(counter.max, expected_max)
    62      self.assertEqual(histogram.firstBucketOffset, expected_first_bucket_index)
    63      self.assertEqual(histogram.bucketCounts, expected_buckets)
    64  
    65    def test_translate_to_histogram_with_input_0(self):
    66      counter = DataflowDistributionCounter()
    67      counter.add_input(0)
    68      histogram = Mock(firstBucketOffset=None, bucketCounts=None)
    69      counter.translate_to_histogram(histogram)
    70      self.assertEqual(histogram.firstBucketOffset, 0)
    71      self.assertEqual(histogram.bucketCounts, [1])
    72  
    73    def test_translate_to_histogram_with_max_input(self):
    74      counter = DataflowDistributionCounter()
    75      counter.add_input(INT64_MAX)
    76      histogram = Mock(firstBucketOffset=None, bucketCounts=None)
    77      counter.translate_to_histogram(histogram)
    78      self.assertEqual(histogram.firstBucketOffset, 57)
    79      self.assertEqual(histogram.bucketCounts, [1])
    80  
    81  
    82  if __name__ == '__main__':
    83    unittest.main()