|
32 | 32 | from apache_beam.metrics.metric import MetricResults |
33 | 33 | from apache_beam.metrics.metric import Metrics |
34 | 34 | from apache_beam.metrics.metric import MetricsFilter |
| 35 | +from apache_beam.metrics.metric import MetricsFlag |
35 | 36 | from apache_beam.metrics.metricbase import MetricName |
| 37 | +from apache_beam.options.pipeline_options import PipelineOptions |
36 | 38 | from apache_beam.runners.direct.direct_runner import BundleBasedDirectRunner |
37 | 39 | from apache_beam.runners.worker import statesampler |
38 | 40 | from apache_beam.testing.metric_result_matchers import DistributionMatcher |
@@ -121,6 +123,108 @@ def test_get_namespace_error(self): |
121 | 123 | with self.assertRaises(ValueError): |
122 | 124 | Metrics.get_namespace(object()) |
123 | 125 |
|
| 126 | + def test_metrics_flag(self): |
| 127 | + MetricsFlag.reset() |
| 128 | + try: |
| 129 | + self.assertFalse(MetricsFlag.counter_disabled) |
| 130 | + self.assertFalse(MetricsFlag.string_set_disabled) |
| 131 | + self.assertFalse(MetricsFlag.bounded_trie_disabled) |
| 132 | + |
| 133 | + options = PipelineOptions(['--experiments=disableCounterMetrics']) |
| 134 | + MetricsFlag.set_default_pipeline_options(options) |
| 135 | + self.assertTrue(MetricsFlag.counter_disabled) |
| 136 | + self.assertFalse(MetricsFlag.string_set_disabled) |
| 137 | + self.assertFalse(MetricsFlag.bounded_trie_disabled) |
| 138 | + |
| 139 | + MetricsFlag.reset() |
| 140 | + options = PipelineOptions(['--experiments=disableStringSetMetrics']) |
| 141 | + MetricsFlag.set_default_pipeline_options(options) |
| 142 | + self.assertFalse(MetricsFlag.counter_disabled) |
| 143 | + self.assertTrue(MetricsFlag.string_set_disabled) |
| 144 | + self.assertFalse(MetricsFlag.bounded_trie_disabled) |
| 145 | + |
| 146 | + MetricsFlag.reset() |
| 147 | + options = PipelineOptions(['--experiments=disableBoundedTrieMetrics']) |
| 148 | + MetricsFlag.set_default_pipeline_options(options) |
| 149 | + self.assertFalse(MetricsFlag.counter_disabled) |
| 150 | + self.assertFalse(MetricsFlag.string_set_disabled) |
| 151 | + self.assertTrue(MetricsFlag.bounded_trie_disabled) |
| 152 | + |
| 153 | + MetricsFlag.reset() |
| 154 | + options = PipelineOptions([ |
| 155 | + '--experiments=disableCounterMetrics', |
| 156 | + '--experiments=disableStringSetMetrics', |
| 157 | + '--experiments=disableBoundedTrieMetrics', |
| 158 | + ]) |
| 159 | + MetricsFlag.set_default_pipeline_options(options) |
| 160 | + self.assertTrue(MetricsFlag.counter_disabled) |
| 161 | + self.assertTrue(MetricsFlag.string_set_disabled) |
| 162 | + self.assertTrue(MetricsFlag.bounded_trie_disabled) |
| 163 | + finally: |
| 164 | + MetricsFlag.reset() |
| 165 | + |
| 166 | + def test_disabled_counter_is_noop(self): |
| 167 | + sampler = statesampler.StateSampler('', counters.CounterFactory()) |
| 168 | + statesampler.set_current_tracker(sampler) |
| 169 | + state = sampler.scoped_state( |
| 170 | + 'mystep', 'myState', metrics_container=MetricsContainer('mystep')) |
| 171 | + MetricsFlag.reset() |
| 172 | + try: |
| 173 | + sampler.start() |
| 174 | + with state: |
| 175 | + container = MetricsEnvironment.current_container() |
| 176 | + Metrics.counter('ns', 'baseline').inc() |
| 177 | + self.assertEqual(len(container.metrics), 1) |
| 178 | + options = PipelineOptions(['--experiments=disableCounterMetrics']) |
| 179 | + MetricsFlag.set_default_pipeline_options(options) |
| 180 | + Metrics.counter('ns', 'after_disable').inc() |
| 181 | + Metrics.counter('ns', 'after_disable').inc(5) |
| 182 | + Metrics.counter('ns', 'after_disable').dec() |
| 183 | + self.assertEqual(len(container.metrics), 1) |
| 184 | + finally: |
| 185 | + sampler.stop() |
| 186 | + MetricsFlag.reset() |
| 187 | + |
| 188 | + def test_disabled_string_set_is_noop(self): |
| 189 | + sampler = statesampler.StateSampler('', counters.CounterFactory()) |
| 190 | + statesampler.set_current_tracker(sampler) |
| 191 | + state = sampler.scoped_state( |
| 192 | + 'mystep', 'myState', metrics_container=MetricsContainer('mystep')) |
| 193 | + MetricsFlag.reset() |
| 194 | + try: |
| 195 | + sampler.start() |
| 196 | + with state: |
| 197 | + container = MetricsEnvironment.current_container() |
| 198 | + Metrics.string_set('ns', 'baseline').add('seed') |
| 199 | + self.assertEqual(len(container.metrics), 1) |
| 200 | + options = PipelineOptions(['--experiments=disableStringSetMetrics']) |
| 201 | + MetricsFlag.set_default_pipeline_options(options) |
| 202 | + Metrics.string_set('ns', 'after_disable').add('value') |
| 203 | + self.assertEqual(len(container.metrics), 1) |
| 204 | + finally: |
| 205 | + sampler.stop() |
| 206 | + MetricsFlag.reset() |
| 207 | + |
| 208 | + def test_disabled_bounded_trie_is_noop(self): |
| 209 | + sampler = statesampler.StateSampler('', counters.CounterFactory()) |
| 210 | + statesampler.set_current_tracker(sampler) |
| 211 | + state = sampler.scoped_state( |
| 212 | + 'mystep', 'myState', metrics_container=MetricsContainer('mystep')) |
| 213 | + MetricsFlag.reset() |
| 214 | + try: |
| 215 | + sampler.start() |
| 216 | + with state: |
| 217 | + container = MetricsEnvironment.current_container() |
| 218 | + Metrics.bounded_trie('ns', 'baseline').add(['a']) |
| 219 | + self.assertEqual(len(container.metrics), 1) |
| 220 | + options = PipelineOptions(['--experiments=disableBoundedTrieMetrics']) |
| 221 | + MetricsFlag.set_default_pipeline_options(options) |
| 222 | + Metrics.bounded_trie('ns', 'after_disable').add(['a', 'b']) |
| 223 | + self.assertEqual(len(container.metrics), 1) |
| 224 | + finally: |
| 225 | + sampler.stop() |
| 226 | + MetricsFlag.reset() |
| 227 | + |
124 | 228 | def test_counter_empty_name(self): |
125 | 229 | with self.assertRaises(ValueError): |
126 | 230 | Metrics.counter("namespace", "") |
|
0 commit comments