| import unittest | |
| from unittest.mock import patch, MagicMock, mock_open | |
| import os | |
| import json | |
| import tempfile | |
| import base64 | |
| import collect_metrics | |
| def create_mock_bep_content(test_results=None, exit_code=0): | |
| content = [] | |
| # 1. Optional Test Results | |
| if test_results: | |
| for tr in test_results: | |
| content.append(json.dumps({ | |
| "id": {"testResult": {"label": tr["label"]}}, | |
| "testResult": {"status": tr["status"], "testAttemptDurationMillis": str(tr["duration_ms"])}, | |
| })) | |
| # 2. Common Build Metrics | |
| content.append(json.dumps({ | |
| "id": {"buildMetrics": {}}, | |
| "buildMetrics": { | |
| "timingMetrics": {"wallTimeInMs": "10000"}, | |
| "actionSummary": { | |
| "actionsExecuted": "100", | |
| "runnerCount": [{"name": "remote cache hit", "count": "50"}], | |
| }, | |
| "artifactMetrics": {"topLevelArtifacts": {"sizeInBytes": "1024"}}, | |
| "networkMetrics": {"systemNetworkStats": {"bytesRecv": "512"}}, | |
| }, | |
| })) | |
| # 3. Common Build Tool Logs (Critical Path) | |
| content.append(json.dumps({ | |
| "id": {"buildToolLogs": {}}, | |
| "buildToolLogs": { | |
| "log": [{ | |
| "name": "critical path", | |
| "contents": base64.b64encode(b"Critical Path: 15.0s\n").decode("utf-8"), | |
| }] | |
| }, | |
| })) | |
| # 4. Common Build Finished | |
| content.append(json.dumps({ | |
| "id": {"buildFinished": {}}, | |
| "finished": {"exitCode": {"code": exit_code}} | |
| })) | |
| return content | |
| class TestPublishMetrics(unittest.TestCase): | |
| def setUp(self): | |
| # Reset environment variables before each test | |
| self.original_environ = os.environ.copy() | |
| os.environ.clear() | |
| def tearDown(self): | |
| os.environ.clear() | |
| os.environ.update(self.original_environ) | |
| # --- Test 1: Git Stats Logic --- | |
| @patch("subprocess.check_output") | |
| def test_get_git_stats_success(self, mock_subprocess): | |
| # Mock successful git output | |
| mock_subprocess.return_value = " 5 files changed, 20 insertions(+), 5 deletions(-)" | |
| count = collect_metrics.get_git_stats() | |
| self.assertEqual(count, 5) | |
| @patch("subprocess.check_output") | |
| def test_get_git_stats_singular(self, mock_subprocess): | |
| # Mock singular output | |
| mock_subprocess.return_value = " 1 file changed, 1 insertion(+)" | |
| count = collect_metrics.get_git_stats() | |
| self.assertEqual(count, 1) | |
| @patch("subprocess.check_output") | |
| def test_get_git_stats_failure(self, mock_subprocess): | |
| # Mock a git failure | |
| mock_subprocess.side_effect = Exception("Git command not found") | |
| count = collect_metrics.get_git_stats() | |
| self.assertEqual(count, 9999) # Fallback value | |
| # --- Test 2: BEP Parsing --- | |
| def test_parse_bep_valid(self): | |
| # Create a mock BEP file content | |
| test_results = [ | |
| {"label": "//pkg:test1", "status": "PASSED", "duration_ms": 1500}, | |
| {"label": "//pkg:test2", "status": "FAILED", "duration_ms": 5000}, | |
| ] | |
| mock_bep_content = create_mock_bep_content(test_results=test_results) | |
| # Create a temp file to hold the data | |
| with tempfile.NamedTemporaryFile(mode="w", delete=False) as tf: | |
| tf.write("\n".join(mock_bep_content)) | |
| temp_path = tf.name | |
| try: | |
| bep_metrics = collect_metrics.parse_bep(temp_path) | |
| # Verify Metrics | |
| self.assertEqual(bep_metrics.wall_time_ms, 10000) | |
| self.assertEqual(bep_metrics.total_actions, 100) | |
| self.assertEqual(bep_metrics.remote_and_disk_cache_hits, 50) | |
| self.assertEqual(bep_metrics.failed_test_count, 1) | |
| self.assertEqual(bep_metrics.critical_path_s, 15.0) | |
| self.assertEqual(bep_metrics.exit_code, 0) | |
| # Verify Targets | |
| self.assertEqual(len(bep_metrics.targets), 2) | |
| target1 = next(t for t in bep_metrics.targets if t.label == "//pkg:test1") | |
| self.assertEqual(target1.status, "PASSED") | |
| self.assertEqual(target1.duration_s, 1.5) | |
| finally: | |
| if os.path.exists(temp_path): | |
| os.remove(temp_path) | |
| # --- Second Run to fall back to outputArtifactsSeen --- | |
| mock_bep_content_2 = [ | |
| json.dumps( | |
| { | |
| "id": {"buildMetrics": {}}, | |
| "buildMetrics": { | |
| "artifactMetrics": { | |
| "topLevelArtifacts": {"sizeInBytes": "0"}, | |
| "outputArtifactsSeen": {"sizeInBytes": "4096"} | |
| }, | |
| }, | |
| } | |
| ), | |
| ] | |
| # Create a temp file to hold the data | |
| with tempfile.NamedTemporaryFile(mode="w", delete=False) as tf2: | |
| tf2.write("\n".join(mock_bep_content_2)) | |
| temp_path_2 = tf2.name | |
| try: | |
| bep_metrics_2 = collect_metrics.parse_bep(temp_path_2) | |
| self.assertEqual(bep_metrics_2.output_size_bytes, 4096) | |
| finally: | |
| if os.path.exists(temp_path_2): | |
| os.remove(temp_path_2) | |
| def test_parse_bep_build_only(self): | |
| # Create a mock BEP file content with only build metrics | |
| mock_bep_content = create_mock_bep_content() | |
| # Create a temp file to hold the data | |
| with tempfile.NamedTemporaryFile(mode="w", delete=False) as tf: | |
| tf.write("\n".join(mock_bep_content)) | |
| temp_path = tf.name | |
| try: | |
| bep_metrics = collect_metrics.parse_bep(temp_path) | |
| # Verify Metrics | |
| self.assertEqual(bep_metrics.wall_time_ms, 10000) | |
| self.assertEqual(bep_metrics.total_actions, 100) | |
| self.assertEqual(bep_metrics.remote_and_disk_cache_hits, 50) | |
| self.assertEqual(bep_metrics.failed_test_count, 0) | |
| self.assertEqual(bep_metrics.critical_path_s, 15.0) | |
| self.assertEqual(bep_metrics.exit_code, 0) | |
| # Verify Targets are empty | |
| self.assertEqual(len(bep_metrics.targets), 0) | |
| finally: | |
| if os.path.exists(temp_path): | |
| os.remove(temp_path) | |
| def test_extract_critical_path(self): | |
| # Mock a Base64 encoded critical path log | |
| raw_log = "Critical Path: 12.5s\n Action A..." | |
| b64_log = base64.b64encode(raw_log.encode("utf-8")).decode("utf-8") | |
| logs = [{"name": "critical path", "contents": b64_log}] | |
| duration = collect_metrics.extract_critical_path(logs) | |
| self.assertEqual(duration, 12.5) | |
| # --- Test 3: Main Logic & BigQuery Push --- | |
| @patch("collect_metrics.subprocess.run") | |
| def test_publish_to_bigquery(self, mock_run): | |
| # Mock the subprocess run to succeed (return code 0) | |
| mock_result = MagicMock() | |
| mock_result.returncode = 0 | |
| mock_run.return_value = mock_result | |
| test_row = {"build_number": 123, "pipeline": "test"} | |
| collect_metrics.publish_to_bigquery(test_row) | |
| # Verify it called subprocess.run with bq insert | |
| mock_run.assert_called_once() | |
| call_args = mock_run.call_args[0][0] | |
| self.assertTrue(any("bq" in arg for arg in call_args)) | |
| self.assertIn("insert", call_args) | |
| @patch("collect_metrics.publish_to_bigquery") | |
| @patch("collect_metrics.parse_bep") | |
| @patch("collect_metrics.get_git_stats") | |
| def test_collect_metrics_end_to_end(self, mock_git, mock_parse, mock_publish): | |
| # Setup Environment | |
| os.environ["BUILDKITE_BUILD_NUMBER"] = "500" | |
| os.environ["BUILDKITE_PIPELINE_SLUG"] = "test-pipeline" | |
| os.environ["BUILDKITE_ORGANIZATION_SLUG"] = "test-org" | |
| # Setup Mocks | |
| mock_git.return_value = 5 # 5 changed files | |
| # Mock BEP Return | |
| mock_bep_metrics = collect_metrics.BazelMetrics( | |
| wall_time_ms=5000, | |
| critical_path_s=4.0, | |
| remote_and_disk_cache_hits=10, | |
| total_actions=20, | |
| output_size_bytes=100, | |
| bytes_downloaded=50, | |
| failed_test_count=0, | |
| exit_code=0, | |
| ) | |
| mock_parse.return_value = mock_bep_metrics | |
| # Run Function (with mocked timestamps) | |
| with patch("collect_metrics.fetch_job_timestamps") as mock_fetch: | |
| mock_fetch.return_value = collect_metrics.JobTimestamps( | |
| created_at="2023-10-25T10:00:00Z", | |
| started_at="2023-10-25T10:05:00Z" | |
| ) | |
| collect_metrics.collect_metrics_and_push_to_bigquery(test_bep_path="dummy_path.json") | |
| # Verify publish_to_bigquery was called | |
| mock_publish.assert_called_once() | |
| # Inspect the row payload that was generated | |
| row = mock_publish.call_args[0][0] | |
| self.assertEqual(row["build_number"], 500) | |
| self.assertEqual(row["pipeline"], "test-pipeline") | |
| self.assertEqual(row["org"], "test-org") | |
| self.assertEqual(row["changed_files_count"], 5) | |
| self.assertEqual(row["test"]["failed_test_count"], 0) | |
| self.assertEqual(row.get("queue_duration_s"), 300.0) | |
| self.assertIsNone(row["task_label"]) | |
| self.assertEqual(row["build_shard_id"], 0) | |
| self.assertEqual(row["build_shard_count"], 1) | |
| @patch("collect_metrics.publish_to_bigquery") | |
| @patch("collect_metrics.parse_bep") | |
| @patch("collect_metrics.get_git_stats") | |
| def test_collect_metrics_combined(self, mock_git, mock_parse, mock_publish): | |
| # Setup Environment | |
| os.environ["BUILDKITE_BUILD_NUMBER"] = "500" | |
| os.environ["BUILDKITE_PIPELINE_SLUG"] = "test-pipeline" | |
| os.environ["BUILDKITE_ORGANIZATION_SLUG"] = "test-org" | |
| # Setup Mocks | |
| mock_git.return_value = 5 | |
| # Mock BEP Return based on filename | |
| def parse_bep_side_effect(filepath): | |
| if "build" in filepath: | |
| return collect_metrics.BazelMetrics( | |
| wall_time_ms=5000, | |
| critical_path_s=4.0, | |
| total_actions=20, | |
| exit_code=0, | |
| ) | |
| elif "test" in filepath: | |
| return collect_metrics.BazelMetrics( | |
| wall_time_ms=10000, | |
| critical_path_s=8.0, | |
| total_actions=50, | |
| exit_code=0, | |
| failed_test_count=1, | |
| targets=[collect_metrics.TestTarget(label="//pkg:test1", status="FAILED", duration_s=5.0)], | |
| ) | |
| return None | |
| mock_parse.side_effect = parse_bep_side_effect | |
| # Run Function | |
| with patch("collect_metrics.fetch_job_timestamps") as mock_fetch: | |
| mock_fetch.return_value = collect_metrics.JobTimestamps( | |
| created_at="2023-10-25T10:00:00Z", | |
| started_at="2023-10-25T10:05:00Z", | |
| finished_at="2023-10-25T10:15:00Z" | |
| ) | |
| collect_metrics.collect_metrics_and_push_to_bigquery( | |
| build_bep_path="build_bep.json", test_bep_path="test_bep.json" | |
| ) | |
| # Verify publish_to_bigquery was called | |
| mock_publish.assert_called_once() | |
| # Inspect the row payload | |
| row = mock_publish.call_args[0][0] | |
| self.assertEqual(row["build_number"], 500) | |
| self.assertIn("build", row) | |
| self.assertIn("test", row) | |
| self.assertEqual(row["build"]["wall_time_s"], 5.0) | |
| self.assertEqual(row["test"]["wall_time_s"], 10.0) | |
| self.assertEqual(row["test"]["failed_test_count"], 1) | |
| self.assertEqual(len(row["test"]["targets"]), 1) | |
| self.assertEqual(row["test"]["targets"][0]["label"], "//pkg:test1") | |
| @patch("collect_metrics.publish_to_bigquery") | |
| @patch("collect_metrics.parse_bep") | |
| @patch("collect_metrics.get_git_stats") | |
| def test_collect_metrics_sharded_task(self, mock_git, mock_parse, mock_publish): | |
| os.environ["BUILDKITE_BUILD_NUMBER"] = "500" | |
| os.environ["BUILDKITE_LABEL"] = "Clang on :ubuntu: Ubuntu 20.04 LTS (shard 1)" | |
| os.environ["BUILDKITE_PARALLEL_JOB"] = "1" | |
| os.environ["BUILDKITE_PARALLEL_JOB_COUNT"] = "3" | |
| mock_parse.return_value = collect_metrics.BazelMetrics(wall_time_ms=5000) | |
| with patch("collect_metrics.fetch_job_timestamps"): | |
| collect_metrics.collect_metrics_and_push_to_bigquery(test_bep_path="dummy.json") | |
| mock_publish.assert_called_once() | |
| row = mock_publish.call_args[0][0] | |
| # Verify that task_label got correctly stripped of the shard suffix and emoji | |
| self.assertEqual(row["task_label"], "Clang on Ubuntu 20.04 LTS") | |
| self.assertEqual(row["build_shard_id"], 1) | |
| self.assertEqual(row["build_shard_count"], 3) | |
| @patch("collect_metrics.subprocess.run") | |
| def test_publish_to_bigquery_failure(self, mock_run): | |
| mock_result = MagicMock() | |
| mock_result.returncode = 1 | |
| mock_result.stdout = "out" | |
| mock_result.stderr = "err" | |
| mock_run.return_value = mock_result | |
| with patch("collect_metrics.print_and_annotate_warning") as mock_annotate: | |
| collect_metrics.publish_to_bigquery({"test": 1}) | |
| mock_annotate.assert_called_once() | |
| def test_duration_parsing_error(self): | |
| with patch.dict(os.environ, {"CHECKOUT_DURATION_S": "invalid", "PREP_DURATION_S": "invalid", "BUILDKITE_BUILD_NUMBER": "500"}): | |
| with patch("collect_metrics.parse_bep") as mock_parse, \ | |
| patch("collect_metrics.publish_to_bigquery"): | |
| mock_parse.return_value = MagicMock() | |
| collect_metrics.collect_metrics_and_push_to_bigquery("dummy") | |
| @patch("collect_metrics.parse_bep") | |
| def test_collect_metrics_bep_failure(self, mock_parse): | |
| mock_parse.return_value = None | |
| with patch.dict(os.environ, {"BUILDKITE_BUILD_NUMBER": "500"}): | |
| with patch("collect_metrics.print_and_annotate_warning") as mock_annotate: | |
| collect_metrics.collect_metrics_and_push_to_bigquery("dummy") | |
| mock_annotate.assert_called_once_with("Skipping BigQuery push due to missing or failed BEP parsing.") | |
| def test_parse_bep_file_not_found(self): | |
| with patch("os.path.exists", return_value=False): | |
| bep_metrics = collect_metrics.parse_bep("non_existent_file.json") | |
| self.assertIsNone(bep_metrics) | |
| @patch("collect_metrics.bazelci.BuildkiteClient") | |
| def test_fetch_job_timestamps_success(self, mock_client_cls): | |
| mock_client = MagicMock() | |
| mock_client_cls.return_value = mock_client | |
| mock_client.get_build_info.return_value = { | |
| "jobs": [{"id": "123", "created_at": "C", "started_at": "S", "finished_at": "F"}] | |
| } | |
| ts = collect_metrics.fetch_job_timestamps("org", "pipe", 1, "123") | |
| self.assertEqual(ts.created_at, "C") | |
| self.assertEqual(ts.finished_at, "F") | |
| if __name__ == "__main__": | |
| unittest.main() |