sdks/python/apache_beam/metrics/execution_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# pytype: skip-file

import functools
import unittest

from apache_beam.metrics.execution import MetricKey
from apache_beam.metrics.execution import MetricsContainer
from apache_beam.metrics.metricbase import MetricName


class TestMetricKey(unittest.TestCase):
  def test_equality_for_key_with_labels(self):
    test_labels = {'label1', 'value1'}
    test_object = MetricKey(
        'step', MetricName('namespace', 'name'), labels=test_labels)
    same_labels = MetricKey(
        'step', MetricName('namespace', 'name'), labels={'label1', 'value1'})
    same_label_reference = MetricKey(
        'step', MetricName('namespace', 'name'), labels=test_labels)
    self.assertEqual(test_object, same_labels)
    self.assertEqual(test_object, same_label_reference)
    self.assertEqual(hash(test_object), hash(same_labels))
    self.assertEqual(hash(test_object), hash(same_label_reference))

  def test_inequality_for_key_with_labels(self):
    test_labels = {'label1', 'value1'}
    test_object = MetricKey(
        'step', MetricName('namespace', 'name'), labels=test_labels)
    no_labels = MetricKey('step', MetricName('namespace', 'name'))
    diff_label_key = MetricKey(
        'step', MetricName('namespace', 'name'), labels={'l1_diff', 'value1'})
    diff_label_value = MetricKey(
        'step', MetricName('namespace', 'name'), labels={'label1', 'v1_diff'})
    self.assertNotEqual(test_object, no_labels)
    self.assertNotEqual(test_object, diff_label_key)
    self.assertNotEqual(test_object, diff_label_value)
    self.assertNotEqual(hash(test_object), hash(no_labels))
    self.assertNotEqual(hash(test_object), hash(diff_label_key))
    self.assertNotEqual(hash(test_object), hash(diff_label_value))

  def test_equality_for_key_with_no_labels(self):
    test_object = MetricKey('step', MetricName('namespace', 'name'))
    same = MetricKey('step', MetricName('namespace', 'name'))
    self.assertEqual(test_object, same)
    self.assertEqual(hash(test_object), hash(same))

    diff_step = MetricKey('step_diff', MetricName('namespace', 'name'))
    diff_namespace = MetricKey('step', MetricName('namespace_diff', 'name'))
    diff_name = MetricKey('step', MetricName('namespace', 'name_diff'))
    self.assertNotEqual(test_object, diff_step)
    self.assertNotEqual(test_object, diff_namespace)
    self.assertNotEqual(test_object, diff_name)
    self.assertNotEqual(hash(test_object), hash(diff_step))
    self.assertNotEqual(hash(test_object), hash(diff_namespace))
    self.assertNotEqual(hash(test_object), hash(diff_name))


class TestMetricsContainer(unittest.TestCase):
  def test_add_to_counter(self):
    mc = MetricsContainer('astep')
    counter = mc.get_counter(MetricName('namespace', 'name'))
    counter.inc()
    counter = mc.get_counter(MetricName('namespace', 'name'))
    self.assertEqual(counter.value, 1)

  def test_get_cumulative_or_updates(self):
    mc = MetricsContainer('astep')

    all_values = []
    for i in range(1, 11):
      counter = mc.get_counter(MetricName('namespace', 'name{}'.format(i)))
      distribution = mc.get_distribution(
          MetricName('namespace', 'name{}'.format(i)))
      gauge = mc.get_gauge(MetricName('namespace', 'name{}'.format(i)))
      str_set = mc.get_string_set(MetricName('namespace', 'name{}'.format(i)))

      counter.inc(i)
      distribution.update(i)
      gauge.set(i)
      str_set.add(str(i % 7))
      all_values.append(i)

    # Retrieve ALL updates.
    cumulative = mc.get_cumulative()
    self.assertEqual(len(cumulative.counters), 10)
    self.assertEqual(len(cumulative.distributions), 10)
    self.assertEqual(len(cumulative.gauges), 10)
    self.assertEqual(len(cumulative.string_sets), 10)

    self.assertEqual(
        set(all_values), {v
                          for _, v in cumulative.counters.items()})
    self.assertEqual(
        set(all_values), {v.value
                          for _, v in cumulative.gauges.items()})
    self.assertEqual(
        {str(i % 7)
         for i in all_values},
        functools.reduce(
            set.union,
            (v.string_set for _, v in cumulative.string_sets.items())))


if __name__ == '__main__':
  unittest.main()