| # |
| # 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. |
| # |
| """ |
| Tests corresponding to the DataflowRunner implementation of MetricsResult, |
| the DataflowMetrics class. |
| """ |
| |
| from __future__ import absolute_import |
| |
| import types |
| import unittest |
| from builtins import object |
| |
| import mock |
| |
| from apache_beam.metrics.cells import DistributionData |
| from apache_beam.metrics.cells import DistributionResult |
| from apache_beam.metrics.execution import MetricKey |
| from apache_beam.metrics.execution import MetricResult |
| from apache_beam.metrics.metricbase import MetricName |
| from apache_beam.runners.dataflow import dataflow_metrics |
| from apache_beam.testing import metric_result_matchers |
| from apache_beam.testing.metric_result_matchers import MetricResultMatcher |
| |
| |
| class DictToObject(object): |
| """Translate from a dict(list()) structure to an object structure""" |
| def __init__(self, data): |
| for name, value in data.items(): |
| setattr(self, name, self._wrap(value)) |
| |
| def _wrap(self, value): |
| if isinstance(value, (tuple, list, set, frozenset)): |
| return type(value)([self._wrap(v) for v in value]) |
| return DictToObject(value) if isinstance(value, dict) else value |
| |
| |
| class TestDataflowMetrics(unittest.TestCase): |
| |
| # TODO(BEAM-6734): Write a dump tool to generate this fake data, or |
| # somehow make this easier to maintain. |
| ONLY_COUNTERS_LIST = {"metrics": [ |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "namespace", |
| "value": "__main__.WordExtractingDoFn"}, |
| {"key": "step", |
| "value": "s2"}, |
| {"key": "tentative", |
| "value": "true"}] |
| }, |
| "name": "words", |
| "origin": "user" |
| }, |
| "scalar": {"integer_value": 26185}, |
| "distribution": None, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "namespace", |
| "value": "__main__.WordExtractingDoFn"}, |
| {"key": "step", |
| "value": "s2"}] |
| }, |
| "name": "words", |
| "origin": "user" |
| }, |
| "scalar": {"integer_value": 26181}, |
| "distribution": None, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "namespace", |
| "value": "__main__.WordExtractingDoFn"}, |
| {"key": "step", |
| "value": "s2"}, |
| {"key": "tentative", |
| "value": "true"}] |
| }, |
| "name": "empty_lines", |
| "origin": "user" |
| }, |
| "scalar": {"integer_value": 1080}, |
| "distribution": None, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "namespace", |
| "value": "__main__.WordExtractingDoFn"}, |
| {"key": "step", |
| "value": "s2"}] |
| }, |
| "name": "empty_lines", |
| "origin": "user" |
| }, |
| "scalar": {"integer_value": 1080}, |
| "distribution": None, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| ]} |
| STRUCTURED_COUNTER_LIST = {"metrics": [ |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "namespace", |
| "value": "__main__.WordExtractingDoFn"}, |
| {"key": "step", |
| "value": "s2"}, |
| {"key": "tentative", |
| "value": "true"}] |
| }, |
| "name": "word_lengths", |
| "origin": "user" |
| }, |
| "scalar": {"integer_value": 109475}, |
| "distribution": None, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "namespace", |
| "value": "__main__.WordExtractingDoFn"}, |
| {"key": "step", |
| "value": "s2"}] |
| }, |
| "name": "word_lengths", |
| "origin": "user" |
| }, |
| "scalar": {"integer_value": 109475}, |
| "distribution": None, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "namespace", |
| "value": "__main__.WordExtractingDoFn"}, |
| {"key": "step", |
| "value": "s2"}, |
| {"key": "tentative", |
| "value": "true"}] |
| }, |
| "name": "word_length_dist", |
| "origin": "user" |
| }, |
| "scalar": None, |
| "distribution": { |
| "object_value": { |
| "properties": [ |
| {"key": "min", "value": |
| {"integer_value": 2}}, |
| {"key": "max", "value": |
| {"integer_value": 16}}, |
| {"key": "count", "value": |
| {"integer_value": 2}}, |
| {"key": "mean", "value": |
| {"integer_value": 9}}, |
| {"key": "sum", "value": |
| {"integer_value": 18}},] |
| } |
| }, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "namespace", |
| "value": "__main__.WordExtractingDoFn"}, |
| {"key": "step", |
| "value": "s2"}] |
| }, |
| "name": "word_length_dist", |
| "origin": "user" |
| }, |
| "scalar": None, |
| "distribution": { |
| "object_value": { |
| "properties": [ |
| {"key": "min", "value": |
| {"integer_value": 2}}, |
| {"key": "max", "value": |
| {"integer_value": 16}}, |
| {"key": "count", "value": |
| {"integer_value": 2}}, |
| {"key": "mean", "value": |
| {"integer_value": 9}}, |
| {"key": "sum", "value": |
| {"integer_value": 18}}, |
| ] |
| } |
| }, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| ]} |
| SYSTEM_COUNTERS_LIST = {"metrics": [ |
| # ElementCount |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "original_name", |
| "value": "ToIsmRecordForMultimap-out0-ElementCount"}, |
| {"key": "output_user_name", |
| "value": "ToIsmRecordForMultimap-out0"} |
| ] |
| }, |
| "name": "ElementCount", |
| "origin": "dataflow/v1b3" |
| }, |
| "scalar": {"integer_value": 42}, |
| "distribution": None, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "original_name", |
| "value": "ToIsmRecordForMultimap-out0-ElementCount"}, |
| {"key": "output_user_name", |
| "value": "ToIsmRecordForMultimap-out0"}, |
| {"key": "tentative", |
| "value": "true"} |
| ] |
| }, |
| "name": "ElementCount", |
| "origin": "dataflow/v1b3" |
| }, |
| "scalar": {"integer_value": 42}, |
| "distribution": None, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| # MeanByteCount |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "original_name", |
| "value": "Read-out0-MeanByteCount"}, |
| {"key": "output_user_name", |
| "value": "GroupByKey/Read-out0"} |
| ] |
| }, |
| "name": "MeanByteCount", |
| "origin": "dataflow/v1b3" |
| }, |
| "scalar": {"integer_value": 31}, |
| "distribution": None, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "original_name", |
| "value": "Read-out0-MeanByteCount"}, |
| {"key": "output_user_name", |
| "value": "GroupByKey/Read-out0"}, |
| {"key": "tentative", |
| "value": "true"} |
| ] |
| }, |
| "name": "MeanByteCount", |
| "origin": "dataflow/v1b3" |
| }, |
| "scalar": {"integer_value": 31}, |
| "distribution": None, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| # ExecutionTime |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "step", |
| "value": "write/Write/Write"}, |
| ] |
| }, |
| "name": "ExecutionTime_ProcessElement", |
| "origin": "dataflow/v1b3" |
| }, |
| "scalar": {"integer_value": 1000}, |
| "distribution": None, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| {"name": {"context": |
| {"additionalProperties": [ |
| {"key": "step", |
| "value": "write/Write/Write"}, |
| {"key": "tentative", |
| "value": "true"} |
| ] |
| }, |
| "name": "ExecutionTime_ProcessElement", |
| "origin": "dataflow/v1b3" |
| }, |
| "scalar": {"integer_value": 1000}, |
| "distribution": None, |
| "updateTime": "2017-03-22T18:47:06.402Z" |
| }, |
| ]} |
| |
| def setup_mock_client_result(self, counter_list=None): |
| mock_client = mock.Mock() |
| mock_query_result = DictToObject(counter_list) |
| mock_client.get_job_metrics.return_value = mock_query_result |
| mock_job_result = mock.Mock() |
| mock_job_result.job_id.return_value = 1 |
| mock_job_result.is_in_terminal_state.return_value = False |
| return mock_client, mock_job_result |
| |
| def test_cache_functions(self): |
| mock_client, mock_job_result = self.setup_mock_client_result( |
| self.STRUCTURED_COUNTER_LIST) |
| dm = dataflow_metrics.DataflowMetrics(mock_client, mock_job_result) |
| |
| # At first creation, we should always query dataflow. |
| self.assertTrue(dm._cached_metrics is None) |
| |
| # Right after querying, we still query again. |
| dm.query() |
| self.assertTrue(dm._cached_metrics is None) |
| |
| # The job has ended. The query should not run again after this. |
| mock_job_result.is_in_terminal_state.return_value = True |
| dm.query() |
| self.assertTrue(dm._cached_metrics) |
| |
| def test_query_structured_metrics(self): |
| mock_client, mock_job_result = self.setup_mock_client_result( |
| self.STRUCTURED_COUNTER_LIST) |
| dm = dataflow_metrics.DataflowMetrics(mock_client, mock_job_result) |
| dm._translate_step_name = types.MethodType(lambda self, x: 'split', dm) |
| query_result = dm.query() |
| expected_counters = [ |
| MetricResult( |
| MetricKey('split', |
| MetricName('__main__.WordExtractingDoFn', 'word_lengths'), |
| ), |
| 109475, 109475), |
| ] |
| self.assertEqual(query_result['counters'], expected_counters) |
| |
| expected_distributions = [ |
| MetricResult( |
| MetricKey('split', |
| MetricName('__main__.WordExtractingDoFn', |
| 'word_length_dist'), |
| ), |
| DistributionResult(DistributionData( |
| 18, 2, 2, 16)), |
| DistributionResult(DistributionData( |
| 18, 2, 2, 16))), |
| ] |
| self.assertEqual(query_result['distributions'], expected_distributions) |
| |
| def test_query_counters(self): |
| mock_client, mock_job_result = self.setup_mock_client_result( |
| self.ONLY_COUNTERS_LIST) |
| dm = dataflow_metrics.DataflowMetrics(mock_client, mock_job_result) |
| dm._translate_step_name = types.MethodType(lambda self, x: 'split', dm) |
| query_result = dm.query() |
| expected_counters = [ |
| MetricResult( |
| MetricKey('split', |
| MetricName('__main__.WordExtractingDoFn', 'empty_lines') |
| ), |
| 1080, 1080), |
| MetricResult( |
| MetricKey('split', |
| MetricName('__main__.WordExtractingDoFn', 'words')), |
| 26181, 26185), |
| ] |
| self.assertEqual(sorted(query_result['counters'], |
| key=lambda x: x.key.metric.name), |
| sorted(expected_counters, |
| key=lambda x: x.key.metric.name)) |
| |
| def test_system_counters_set_labels_and_step_name(self): |
| mock_client, mock_job_result = self.setup_mock_client_result( |
| self.SYSTEM_COUNTERS_LIST) |
| test_object = dataflow_metrics.DataflowMetrics(mock_client, mock_job_result) |
| all_metrics = test_object.all_metrics() |
| |
| matchers = [ |
| MetricResultMatcher( |
| name='ElementCount', |
| labels={ |
| 'original_name' : 'ToIsmRecordForMultimap-out0-ElementCount', |
| 'output_user_name' : 'ToIsmRecordForMultimap-out0' |
| }, |
| attempted=42, |
| committed=42 |
| ), |
| MetricResultMatcher( |
| name='MeanByteCount', |
| labels={ |
| 'original_name' : 'Read-out0-MeanByteCount', |
| 'output_user_name' : 'GroupByKey/Read-out0' |
| }, |
| attempted=31, |
| committed=31 |
| ), |
| MetricResultMatcher( |
| name='ExecutionTime_ProcessElement', |
| step='write/Write/Write', |
| attempted=1000, |
| committed=1000 |
| ) |
| ] |
| errors = metric_result_matchers.verify_all(all_metrics, matchers) |
| self.assertFalse(errors, errors) |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
