github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/runners/direct/direct_metrics_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 import hamcrest as hc 23 24 from apache_beam.metrics.cells import DistributionData 25 from apache_beam.metrics.cells import DistributionResult 26 from apache_beam.metrics.execution import MetricKey 27 from apache_beam.metrics.execution import MetricResult 28 from apache_beam.metrics.execution import MetricUpdates 29 from apache_beam.metrics.metricbase import MetricName 30 from apache_beam.runners.direct.direct_metrics import DirectMetrics 31 32 33 class DirectMetricsTest(unittest.TestCase): 34 name1 = MetricName('namespace1', 'name1') 35 name2 = MetricName('namespace1', 'name2') 36 name3 = MetricName('namespace2', 'name1') 37 38 bundle1 = object() # For this test, any object can be a bundle 39 bundle2 = object() 40 41 def test_combiner_functions(self): 42 metrics = DirectMetrics() 43 counter = metrics._counters['anykey'] 44 counter.commit_logical(self.bundle1, 5) 45 self.assertEqual(counter.extract_committed(), 5) 46 with self.assertRaises(TypeError): 47 counter.commit_logical(self.bundle1, None) 48 49 distribution = metrics._distributions['anykey'] 50 distribution.commit_logical(self.bundle1, DistributionData(4, 1, 4, 4)) 51 self.assertEqual( 52 distribution.extract_committed(), 53 DistributionResult(DistributionData(4, 1, 4, 4))) 54 55 with self.assertRaises(AttributeError): 56 distribution.commit_logical(self.bundle1, None) 57 58 def test_commit_logical_no_filter(self): 59 metrics = DirectMetrics() 60 metrics.commit_logical( 61 self.bundle1, 62 MetricUpdates( 63 counters={ 64 MetricKey('step1', self.name1): 5, 65 MetricKey('step1', self.name2): 8 66 }, 67 distributions={ 68 MetricKey('step1', self.name1): DistributionData(8, 2, 3, 5) 69 })) 70 71 metrics.commit_logical( 72 self.bundle1, 73 MetricUpdates( 74 counters={ 75 MetricKey('step2', self.name1): 7, 76 MetricKey('step1', self.name2): 4 77 }, 78 distributions={ 79 MetricKey('step1', self.name1): DistributionData(4, 1, 4, 4) 80 })) 81 82 results = metrics.query() 83 hc.assert_that( 84 results['counters'], 85 hc.contains_inanyorder( 86 *[ 87 MetricResult(MetricKey('step1', self.name2), 12, 0), 88 MetricResult(MetricKey('step2', self.name1), 7, 0), 89 MetricResult(MetricKey('step1', self.name1), 5, 0) 90 ])) 91 hc.assert_that( 92 results['distributions'], 93 hc.contains_inanyorder( 94 MetricResult( 95 MetricKey('step1', self.name1), 96 DistributionResult(DistributionData(12, 3, 3, 5)), 97 DistributionResult(DistributionData(0, 0, None, None))))) 98 99 def test_apply_physical_no_filter(self): 100 metrics = DirectMetrics() 101 metrics.update_physical( 102 object(), 103 MetricUpdates( 104 counters={ 105 MetricKey('step1', self.name1): 5, 106 MetricKey('step1', self.name3): 8 107 })) 108 109 metrics.update_physical( 110 object(), 111 MetricUpdates( 112 counters={ 113 MetricKey('step2', self.name1): 7, 114 MetricKey('step1', self.name3): 4 115 })) 116 results = metrics.query() 117 hc.assert_that( 118 results['counters'], 119 hc.contains_inanyorder( 120 *[ 121 MetricResult(MetricKey('step1', self.name1), 0, 5), 122 MetricResult(MetricKey('step1', self.name3), 0, 12), 123 MetricResult(MetricKey('step2', self.name1), 0, 7) 124 ])) 125 126 metrics.commit_physical(object(), MetricUpdates()) 127 results = metrics.query() 128 hc.assert_that( 129 results['counters'], 130 hc.contains_inanyorder( 131 *[ 132 MetricResult(MetricKey('step1', self.name1), 0, 5), 133 MetricResult(MetricKey('step1', self.name3), 0, 12), 134 MetricResult(MetricKey('step2', self.name1), 0, 7) 135 ])) 136 137 def test_apply_physical_logical(self): 138 metrics = DirectMetrics() 139 dist_zero = DistributionData(0, 0, None, None) 140 metrics.update_physical( 141 object(), 142 MetricUpdates( 143 counters={ 144 MetricKey('step1', self.name1): 7, 145 MetricKey('step1', self.name2): 5, 146 MetricKey('step2', self.name1): 1 147 }, 148 distributions={ 149 MetricKey('step1', self.name1): DistributionData(3, 1, 3, 3), 150 MetricKey('step2', self.name3): DistributionData(8, 2, 4, 4) 151 })) 152 results = metrics.query() 153 hc.assert_that( 154 results['counters'], 155 hc.contains_inanyorder( 156 *[ 157 MetricResult(MetricKey('step1', self.name1), 0, 7), 158 MetricResult(MetricKey('step1', self.name2), 0, 5), 159 MetricResult(MetricKey('step2', self.name1), 0, 1) 160 ])) 161 hc.assert_that( 162 results['distributions'], 163 hc.contains_inanyorder( 164 *[ 165 MetricResult( 166 MetricKey('step1', self.name1), 167 DistributionResult(dist_zero), 168 DistributionResult(DistributionData(3, 1, 3, 3))), 169 MetricResult( 170 MetricKey('step2', self.name3), 171 DistributionResult(dist_zero), 172 DistributionResult(DistributionData(8, 2, 4, 4))) 173 ])) 174 175 metrics.commit_physical( 176 object(), 177 MetricUpdates( 178 counters={ 179 MetricKey('step1', self.name1): -3, 180 MetricKey('step2', self.name1): -5 181 }, 182 distributions={ 183 MetricKey('step1', self.name1): DistributionData(8, 4, 1, 5), 184 MetricKey('step2', self.name2): DistributionData(8, 8, 1, 1) 185 })) 186 results = metrics.query() 187 hc.assert_that( 188 results['counters'], 189 hc.contains_inanyorder( 190 *[ 191 MetricResult(MetricKey('step1', self.name1), 0, 4), 192 MetricResult(MetricKey('step1', self.name2), 0, 5), 193 MetricResult(MetricKey('step2', self.name1), 0, -4) 194 ])) 195 hc.assert_that( 196 results['distributions'], 197 hc.contains_inanyorder( 198 *[ 199 MetricResult( 200 MetricKey('step1', self.name1), 201 DistributionResult(dist_zero), 202 DistributionResult(DistributionData(11, 5, 1, 5))), 203 MetricResult( 204 MetricKey('step2', self.name3), 205 DistributionResult(dist_zero), 206 DistributionResult(DistributionData(8, 2, 4, 4))), 207 MetricResult( 208 MetricKey('step2', self.name2), 209 DistributionResult(dist_zero), 210 DistributionResult(DistributionData(8, 8, 1, 1))) 211 ])) 212 213 metrics.commit_logical( 214 object(), 215 MetricUpdates( 216 counters={ 217 MetricKey('step1', self.name1): 3, 218 MetricKey('step1', self.name2): 5, 219 MetricKey('step2', self.name1): -3 220 }, 221 distributions={ 222 MetricKey('step1', self.name1): DistributionData(11, 5, 1, 5), 223 MetricKey('step2', self.name2): DistributionData(8, 8, 1, 1), 224 MetricKey('step2', self.name3): DistributionData(4, 1, 4, 4) 225 })) 226 227 results = metrics.query() 228 hc.assert_that( 229 results['counters'], 230 hc.contains_inanyorder( 231 *[ 232 MetricResult(MetricKey('step1', self.name1), 3, 4), 233 MetricResult(MetricKey('step1', self.name2), 5, 5), 234 MetricResult(MetricKey('step2', self.name1), -3, -4) 235 ])) 236 hc.assert_that( 237 results['distributions'], 238 hc.contains_inanyorder( 239 *[ 240 MetricResult( 241 MetricKey('step1', self.name1), 242 DistributionResult(DistributionData(11, 5, 1, 5)), 243 DistributionResult(DistributionData(11, 5, 1, 5))), 244 MetricResult( 245 MetricKey('step2', self.name3), 246 DistributionResult(DistributionData(4, 1, 4, 4)), 247 DistributionResult(DistributionData(8, 2, 4, 4))), 248 MetricResult( 249 MetricKey('step2', self.name2), 250 DistributionResult(DistributionData(8, 8, 1, 1)), 251 DistributionResult(DistributionData(8, 8, 1, 1))) 252 ])) 253 254 255 if __name__ == '__main__': 256 unittest.main()