Correctly handle aggregation for Azure "count" metrics.

PiperOrigin-RevId: 853327424

Author: bvliu

perfkitbenchmarker/providers/azure/azure_flexible_server.py

@@ -434,6 +434,9 @@ def _CollectProviderMetric(
     logging.info(
         'Collecting metric %s for instance %s', metric_name, self.instance_id
     )
+    aggregation = 'Average'
+    if 'count' in metric_name:
+      aggregation = 'Total'
     cmd = [
         azure.AZURE_PATH,
         'monitor',
@@ -454,7 +457,7 @@ def _CollectProviderMetric(
         '--interval',
         'pt1m',
         '--aggregation',
-        'Average',
+        aggregation,
     ]
     try:
       stdout, _ = vm_util.IssueRetryableCommand(cmd)
@@ -481,10 +484,11 @@ def _CollectProviderMetric(
       return []
 
     points = []
+    key = aggregation.lower()
     for dp in datapoints:
-      if dp['average'] is None:
+      if dp[key] is None:
         continue
-      value = dp['average']
+      value = dp[key]
       if metric.conversion_func:
         value = metric.conversion_func(value)
       points.append((

tests/providers/azure/azure_flexible_server_test.py

@@ -8,6 +8,7 @@
 import mock
 from perfkitbenchmarker import errors
 from perfkitbenchmarker import provider_info
+from perfkitbenchmarker import relational_db
 from perfkitbenchmarker import sql_engine_utils
 from perfkitbenchmarker import vm_util
 from perfkitbenchmarker.providers.azure import azure_network
@@ -113,6 +114,37 @@ def testCollectMetrics(self):
     self.assertEqual(cpu_max.value, 20.0)
     self.assertEqual(cpu_max.unit, '%')
 
+  def testCollectMetricsAggregation(self):
+    mock_issue_cmd = self.enter_context(
+        mock.patch.object(
+            vm_util,
+            'IssueRetryableCommand',
+            return_value=(json.dumps({'value': []}), ''),
+        )
+    )
+    start_time = datetime.datetime(2025, 11, 26, 10, 0, 0)
+    end_time = datetime.datetime(2025, 11, 26, 10, 1, 0)
+
+    # Test 'Total' aggregation
+    count_metric = relational_db.MetricSpec(
+        'some_count', 'some_count', 'count', None
+    )
+    self.server._CollectProviderMetric(count_metric, start_time, end_time)
+
+    call_args = mock_issue_cmd.call_args[0][0]
+    self.assertIn('--aggregation', call_args)
+    self.assertIn('Total', call_args)
+
+    # Test 'Average' aggregation
+    avg_metric = relational_db.MetricSpec(
+        'some_metric', 'some_metric', 'percent', None
+    )
+    self.server._CollectProviderMetric(avg_metric, start_time, end_time)
+
+    call_args = mock_issue_cmd.call_args[0][0]
+    self.assertIn('--aggregation', call_args)
+    self.assertIn('Average', call_args)
+
 
 class AzureFlexibleServerCreateTestCase(pkb_common_test_case.PkbCommonTestCase):