github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/internal/metrics/metric_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 unittest
    21  
    22  from mock import patch
    23  
    24  from apache_beam.internal.metrics.cells import HistogramCellFactory
    25  from apache_beam.internal.metrics.metric import Metrics as InternalMetrics
    26  from apache_beam.internal.metrics.metric import MetricLogger
    27  from apache_beam.metrics.execution import MetricsContainer
    28  from apache_beam.metrics.execution import MetricsEnvironment
    29  from apache_beam.metrics.metric import Metrics
    30  from apache_beam.metrics.metricbase import MetricName
    31  from apache_beam.runners.worker import statesampler
    32  from apache_beam.utils import counters
    33  from apache_beam.utils.histogram import LinearBucket
    34  
    35  
    36  class MetricLoggerTest(unittest.TestCase):
    37    @patch('apache_beam.internal.metrics.metric._LOGGER')
    38    def test_log_metrics(self, mock_logger):
    39      logger = MetricLogger()
    40      logger.minimum_logging_frequency_msec = -1
    41      namespace = Metrics.get_namespace(self.__class__)
    42      metric_name = MetricName(namespace, 'metric_logger_test')
    43      logger.update(HistogramCellFactory(LinearBucket(0, 1, 10)), metric_name, 1)
    44      logger.log_metrics()
    45  
    46      class Contains(str):
    47        def __eq__(self, other):
    48          return self in other
    49  
    50      mock_logger.info.assert_called_once_with(
    51          Contains('HistogramData(Total count: 1, P99: 2, P90: 2, P50: 2)'))
    52  
    53  
    54  class MetricsTest(unittest.TestCase):
    55    def test_create_process_wide(self):
    56      sampler = statesampler.StateSampler('', counters.CounterFactory())
    57      statesampler.set_current_tracker(sampler)
    58      state1 = sampler.scoped_state(
    59          'mystep', 'myState', metrics_container=MetricsContainer('mystep'))
    60  
    61      try:
    62        sampler.start()
    63        with state1:
    64          urn = "my:custom:urn"
    65          labels = {'key': 'value'}
    66          counter = InternalMetrics.counter(
    67              urn=urn, labels=labels, process_wide=True)
    68          # Test that if process_wide is set, that it will be set
    69          # on the process_wide container.
    70          counter.inc(10)
    71          self.assertTrue(isinstance(counter, Metrics.DelegatingCounter))
    72  
    73          del counter
    74  
    75          metric_name = MetricName(None, None, urn=urn, labels=labels)
    76          # Expect a value set on the current container.
    77          self.assertEqual(
    78              MetricsEnvironment.process_wide_container().get_counter(
    79                  metric_name).get_cumulative(),
    80              10)
    81          # Expect no value set on the current container.
    82          self.assertEqual(
    83              MetricsEnvironment.current_container().get_counter(
    84                  metric_name).get_cumulative(),
    85              0)
    86      finally:
    87        sampler.stop()
    88  
    89  
    90  if __name__ == '__main__':
    91    unittest.main()