From c12dfd51952a18afcd094713e9541816b82a5037 Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman@google.com>
Date: Thu, 24 Jul 2025 12:25:31 -0700
Subject: [PATCH] [CI] Add unittesting for metrics collection script

This patch adds some initial unittests for the metrics collection
script. This initial patch focuses on getting the files setup and adding
unittests for uploading metrics. A subsequent patch will add tests for
the workflow collection which is significantly more complicated.

Reviewers: Keenuts, gburgessiv, dschuff, cmtice, lnihlen

Reviewed By: cmtice

Pull Request: https://github.com/llvm/llvm-project/pull/150360
---
 .ci/metrics/metrics_test.py    | 75 ++++++++++++++++++++++++++++++++++
 .github/workflows/check-ci.yml |  1 +
 2 files changed, 76 insertions(+)
 create mode 100644 .ci/metrics/metrics_test.py

diff --git a/.ci/metrics/metrics_test.py b/.ci/metrics/metrics_test.py
new file mode 100644
index 000000000000..259e55f81793
--- /dev/null
+++ b/.ci/metrics/metrics_test.py
@@ -0,0 +1,75 @@
+# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+"""Tests for metrics.py"""
+
+from dataclasses import dataclass
+import requests
+import unittest
+import unittest.mock
+
+import metrics
+
+
+class TestMetrics(unittest.TestCase):
+    def test_upload_gauge_metric(self):
+        """Test that we can upload a gauge metric correctly.
+
+        Also verify that we pass around parameters like API keys and user IDs
+        correctly to the HTTP POST request.
+        """
+        test_metrics = [metrics.GaugeMetric("gauge_test", 5, 1000)]
+        return_value = requests.Response()
+        return_value.status_code = 204
+        with unittest.mock.patch(
+            "requests.post", return_value=return_value
+        ) as post_mock:
+            metrics.upload_metrics(test_metrics, "test_userid", "test_api_key")
+            self.assertSequenceEqual(post_mock.call_args.args, [metrics.GRAFANA_URL])
+            self.assertEqual(
+                post_mock.call_args.kwargs["data"], "gauge_test value=5 1000"
+            )
+            self.assertEqual(
+                post_mock.call_args.kwargs["auth"], ("test_userid", "test_api_key")
+            )
+
+    def test_upload_job_metric(self):
+        """Test that we can upload a job metric correctly."""
+        test_metrics = [
+            metrics.JobMetrics("test_job", 5, 10, 1, 1000, 7, "test_workflow")
+        ]
+        return_value = requests.Response()
+        return_value.status_code = 204
+        with unittest.mock.patch(
+            "requests.post", return_value=return_value
+        ) as post_mock:
+            metrics.upload_metrics(test_metrics, "test_userid", "test_aoi_key")
+            self.assertEqual(
+                post_mock.call_args.kwargs["data"],
+                "test_job queue_time=5,run_time=10,status=1 1000",
+            )
+
+    def test_upload_unknown_metric(self):
+        """Test we report an error if we encounter an unknown metric type."""
+
+        @dataclass
+        class FakeMetric:
+            fake_data: str
+
+        test_metrics = [FakeMetric("test")]
+
+        with self.assertRaises(ValueError):
+            metrics.upload_metrics(test_metrics, "test_userid", "test_api_key")
+
+    def test_bad_response_code(self):
+        """Test that we gracefully handle HTTP response errors."""
+        test_metrics = [metrics.GaugeMetric("gauge_test", 5, 1000)]
+        return_value = requests.Response()
+        return_value.status_code = 403
+        # Just assert that we continue running here and do not raise anything.
+        with unittest.mock.patch("requests.post", return_value=return_value) as _:
+            metrics.upload_metrics(test_metrics, "test_userid", "test_api_key")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/.github/workflows/check-ci.yml b/.github/workflows/check-ci.yml
index d6212c8a8473..befea2093f90 100644
--- a/.github/workflows/check-ci.yml
+++ b/.github/workflows/check-ci.yml
@@ -31,6 +31,7 @@ jobs:
       - name: Install Python Dependencies
         run: |
           pip3 install -r .ci/all_requirements.txt
+          pip3 install -r .ci/metrics/requirements.lock.txt
           pip3 install pytest==8.4.1
       - name: Run Tests
         working-directory: .ci