From 61533450fb92fff63c3adc1adbcdaa4a42575677 Mon Sep 17 00:00:00 2001 From: Divyansh Yadav Date: Tue, 20 Oct 2026 03:03:03 +0530 Subject: [PATCH] Add retries and timeout to Grafana request in flaky test detection --- .test-infra/tools/flaky_test_detection.py | 20 +- .../tools/test_flaky_test_detection.py | 173 ++++++++++++++++++ 2 files changed, 191 insertions(+), 2 deletions(-) create mode 100644 .test-infra/tools/test_flaky_test_detection.py diff --git a/.test-infra/tools/flaky_test_detection.py b/.test-infra/tools/flaky_test_detection.py index de07c63f976d..bd117312aa4f 100644 --- a/.test-infra/tools/flaky_test_detection.py +++ b/.test-infra/tools/flaky_test_detection.py @@ -19,6 +19,8 @@ from datetime import datetime from github import Github from github import Auth +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry ALERT_NAME = "flaky_test" @@ -37,6 +39,8 @@ "go": 10, "misc": 11, } +GRAFANA_REQUEST_TIMEOUT_SECONDS = 30 +GRAFANA_REQUEST_RETRIES = 5 READ_ONLY = os.environ.get("READ_ONLY", "true") @@ -63,7 +67,9 @@ def __init__( def get_workflow_issues(issues): workflows = {} for issue in issues: - for label in issue.get_labels(): + # issue.labels comes with the issue listing, while get_labels() makes + # an extra API request for every issue. + for label in issue.labels: match = re.search(r"workflow_id:\s*(\d+)", str(label.name)) if match: workflow_id = match.group(1) @@ -105,7 +111,17 @@ def create_github_issue(repo, alert): def get_grafana_alerts(): url = f"{GRAFANA_URL}/api/alertmanager/grafana/api/v2/alerts?active=true&filter=alertname%3D{ALERT_NAME}" - response = requests.get(url) + retries = Retry( + total=GRAFANA_REQUEST_RETRIES, + backoff_factor=2, + status_forcelist=[429, 500, 502, 503, 504], + allowed_methods=["GET"], + raise_on_status=False, + ) + with requests.Session() as session: + session.mount("http://", HTTPAdapter(max_retries=retries)) + session.mount("https://", HTTPAdapter(max_retries=retries)) + response = session.get(url, timeout=GRAFANA_REQUEST_TIMEOUT_SECONDS) if response.status_code != 200: raise RuntimeError( "Request to %s failed with status %d: %s" diff --git a/.test-infra/tools/test_flaky_test_detection.py b/.test-infra/tools/test_flaky_test_detection.py new file mode 100644 index 000000000000..11ec77d0c5e8 --- /dev/null +++ b/.test-infra/tools/test_flaky_test_detection.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import json +import threading +import unittest +from datetime import datetime +from http.server import BaseHTTPRequestHandler, HTTPServer +from unittest import mock + +import flaky_test_detection +from flaky_test_detection import get_grafana_alerts, get_workflow_issues + + +ALERT_JSON = [ + { + "labels": { + "workflow_id": "123", + "workflow_url": "https://github.com/apache/beam/actions/workflows/foo.yml", + "workflow_name": "Foo", + "workflow_filename": "foo.yml", + "workflow_threshold": "0.5", + "dashboard_category": "misc", + "workflow_retrieved_at": "2026-01-01T00:00:00+00:00", + } + } +] + + +class FakeGrafanaServer: + """Local HTTP server that replies with the given (status, body) pairs in order.""" + + def __init__(self, responses): + self.responses = list(responses) + self.request_count = 0 + server = self + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + server.request_count += 1 + status, body = server.responses.pop(0) + payload = body.encode("utf-8") + self.send_response(status) + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *args): + pass + + self.httpd = HTTPServer(("127.0.0.1", 0), Handler) + self.url = f"http://127.0.0.1:{self.httpd.server_port}" + + def __enter__(self): + threading.Thread(target=self.httpd.serve_forever, daemon=True).start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.httpd.shutdown() + self.httpd.server_close() + + +class GetGrafanaAlertsTest(unittest.TestCase): + """Tests for get_grafana_alerts.""" + + def fetch_alerts(self, responses): + with FakeGrafanaServer(responses) as server: + with mock.patch.object(flaky_test_detection, "GRAFANA_URL", server.url): + try: + return get_grafana_alerts() + finally: + self.request_count = server.request_count + + def test_parses_alerts(self): + """Test that alerts are built from the Grafana response.""" + alerts = self.fetch_alerts([(200, json.dumps(ALERT_JSON))]) + + self.assertEqual(self.request_count, 1) + self.assertEqual(len(alerts), 1) + alert = alerts[0] + self.assertEqual(alert.workflow_id, "123") + self.assertEqual(alert.workflow_name, "Foo") + self.assertEqual(alert.workflow_filename, "foo.yml") + self.assertEqual(alert.workflow_threshold, 0.5) + self.assertEqual(alert.workflow_dashboard_category, "misc") + self.assertEqual( + alert.workflow_retrieved_at, + datetime.fromisoformat("2026-01-01T00:00:00+00:00"), + ) + + def test_empty_response(self): + """Test that an empty response yields no alerts.""" + alerts = self.fetch_alerts([(200, "")]) + + self.assertEqual(alerts, []) + + def test_retries_on_server_error(self): + """Test that transient server errors are retried.""" + alerts = self.fetch_alerts( + [(503, "unavailable"), (200, json.dumps(ALERT_JSON))] + ) + + self.assertEqual(self.request_count, 2) + self.assertEqual(len(alerts), 1) + + def test_raises_on_client_error(self): + """Test that a non-retryable error fails without retrying.""" + with self.assertRaises(RuntimeError): + self.fetch_alerts([(404, "not found")]) + + self.assertEqual(self.request_count, 1) + + def test_uses_timeout(self): + """Test that the request is made with a timeout.""" + response = mock.Mock(status_code=200, text="") + with mock.patch("requests.Session.get", return_value=response) as get: + get_grafana_alerts() + + self.assertEqual( + get.call_args.kwargs["timeout"], + flaky_test_detection.GRAFANA_REQUEST_TIMEOUT_SECONDS, + ) + + +class GetWorkflowIssuesTest(unittest.TestCase): + """Tests for get_workflow_issues.""" + + @staticmethod + def make_issue(*label_names): + issue = mock.Mock() + issue.labels = [] + for name in label_names: + label = mock.Mock() + label.name = name + issue.labels.append(label) + return issue + + def test_maps_workflow_ids_to_issues(self): + """Test that issues are keyed by the workflow id in their labels.""" + first = self.make_issue("flaky_test", "workflow_id: 111", "bug") + second = self.make_issue("workflow_id:222", "P1") + unrelated = self.make_issue("flaky_test") + + workflows = get_workflow_issues([first, second, unrelated]) + + self.assertEqual(workflows, {"111": first, "222": second}) + + def test_does_not_fetch_labels_per_issue(self): + """Test that labels from the issue listing are used without extra requests.""" + issue = self.make_issue("workflow_id: 111") + + get_workflow_issues([issue]) + + issue.get_labels.assert_not_called() + + +if __name__ == "__main__": + unittest.main()