github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/internal/metrics/cells_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  # pytype: skip-file
    19  
    20  import threading
    21  import unittest
    22  
    23  from apache_beam.internal.metrics.cells import HistogramCell
    24  from apache_beam.internal.metrics.cells import HistogramCellFactory
    25  from apache_beam.internal.metrics.cells import HistogramData
    26  from apache_beam.utils.histogram import Histogram
    27  from apache_beam.utils.histogram import LinearBucket
    28  
    29  
    30  class TestHistogramCell(unittest.TestCase):
    31    @classmethod
    32    def _modify_histogram(cls, d):
    33      for i in range(cls.NUM_ITERATIONS):
    34        d.update(i)
    35  
    36    NUM_THREADS = 5
    37    NUM_ITERATIONS = 100
    38  
    39    def test_parallel_access(self):
    40      # We create NUM_THREADS threads that concurrently modify the distribution.
    41      threads = []
    42      bucket_type = LinearBucket(0, 1, 100)
    43      d = HistogramCell(bucket_type)
    44      for _ in range(TestHistogramCell.NUM_THREADS):
    45        t = threading.Thread(
    46            target=TestHistogramCell._modify_histogram, args=(d, ))
    47        threads.append(t)
    48        t.start()
    49  
    50      for t in threads:
    51        t.join()
    52  
    53      histogram = Histogram(bucket_type)
    54      for _ in range(self.NUM_THREADS):
    55        for i in range(self.NUM_ITERATIONS):
    56          histogram.record(i)
    57  
    58      self.assertEqual(d.get_cumulative(), HistogramData(histogram))
    59  
    60    def test_basic_operations(self):
    61      d = HistogramCellFactory(LinearBucket(0, 1, 10))()
    62      d.update(10)
    63      self.assertEqual(
    64          str(d.get_cumulative()),
    65          'HistogramData(Total count: 1, P99: >=10, P90: >=10, P50: >=10)')
    66      d.update(0)
    67      self.assertEqual(
    68          str(d.get_cumulative()),
    69          'HistogramData(Total count: 2, P99: >=10, P90: >=10, P50: 1)')
    70      d.update(5)
    71      self.assertEqual(
    72          str(d.get_cumulative()),
    73          'HistogramData(Total count: 3, P99: >=10, P90: >=10, P50: 6)')
    74  
    75  
    76  if __name__ == '__main__':
    77    unittest.main()