From 1404fb548111fc796f9a5c688231198691e80e2a Mon Sep 17 00:00:00 2001
From: Deborah Brouwer <deborah.brouwer@collabora.com>
Date: Fri, 13 Dec 2024 17:31:05 -0800
Subject: [PATCH] ci: make the gantt scripts available as modules

Instead of running the gantt scripts only from the command line, make
them available as modules also. This allows the scripts to be imported for
testing.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/32637>
---
 bin/ci/ci_gantt_chart.py | 44 +++++++++++++++++++++-------------------
 bin/ci/ci_post_gantt.py  | 44 ++++++++++++++++++++++------------------
 2 files changed, 47 insertions(+), 41 deletions(-)

diff --git a/bin/ci/ci_gantt_chart.py b/bin/ci/ci_gantt_chart.py
index ddebde55931..36dcf4e808e 100755
--- a/bin/ci/ci_gantt_chart.py
+++ b/bin/ci/ci_gantt_chart.py
@@ -132,7 +132,26 @@ def generate_gantt_chart(pipeline: ProjectPipeline):
     return fig
 
 
-def parse_args() -> argparse.Namespace:
+def main(
+    token: str | None,
+    pipeline_url: str,
+    output: str | None,
+):
+    token = read_token(token)
+    gl = gitlab.Gitlab(url=GITLAB_URL, private_token=token, retry_transient_errors=True)
+
+    pipeline, _ = get_gitlab_pipeline_from_url(gl, pipeline_url)
+    fig = generate_gantt_chart(pipeline)
+    if output and "htm" in output:
+        fig.write_html(output)
+    elif output:
+        fig.update_layout(width=1000)
+        fig.write_image(output)
+    else:
+        fig.show()
+
+
+if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="Generate the Gantt chart from a given pipeline."
     )
@@ -141,29 +160,12 @@ def parse_args() -> argparse.Namespace:
         "-o",
         "--output",
         type=str,
-        help="Output file name. Use html ou image suffixes to choose the format.",
+        help="Output file name. Use html or image suffixes to choose the format.",
     )
     parser.add_argument(
         "--token",
         metavar="token",
         help="force GitLab token, otherwise it's read from ~/.config/gitlab-token",
     )
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-
-    token = read_token(args.token)
-
-    gl = gitlab.Gitlab(url=GITLAB_URL, private_token=token, retry_transient_errors=True)
-
-    pipeline, _ = get_gitlab_pipeline_from_url(gl, args.pipeline_url)
-    fig = generate_gantt_chart(pipeline)
-    if args.output and "htm" in args.output:
-        fig.write_html(args.output)
-    elif args.output:
-        fig.update_layout(width=1000)
-        fig.write_image(args.output)
-    else:
-        fig.show()
+    args = parser.parse_args()
+    main(args.token, args.pipeline_url, args.output)
diff --git a/bin/ci/ci_post_gantt.py b/bin/ci/ci_post_gantt.py
index 04f7939e1cc..70137b75698 100755
--- a/bin/ci/ci_post_gantt.py
+++ b/bin/ci/ci_post_gantt.py
@@ -106,31 +106,15 @@ def gitlab_post_reply_to_note(gl: Gitlab, event: RESTObject, reply_message: str)
         return None
 
 
-def parse_args() -> None:
-    parser = argparse.ArgumentParser(description="Monitor rejected pipelines by Marge.")
-    parser.add_argument(
-        "--token",
-        metavar="token",
-        help="force GitLab token, otherwise it's read from ~/.config/gitlab-token",
-    )
-    parser.add_argument(
-        "--since",
-        metavar="since",
-        help="consider only events after this date (ISO format), otherwise it's read from ~/.config/last_marge_event",
-    )
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
+def main(token: str | None, since: str | None):
     log.basicConfig(level=log.INFO)
 
-    token = read_token(args.token)
+    token = read_token(token)
 
     gl = gitlab.Gitlab(url=GITLAB_URL, private_token=token, retry_transient_errors=True)
 
     user = gl.users.get(MARGE_USER_ID)
-    last_event_at = args.since if args.since else read_last_event_date_from_file()
+    last_event_at = since if since else read_last_event_date_from_file()
 
     log.info(f"Retrieving Marge messages since {pretty_time(last_event_at)}\n")
 
@@ -174,9 +158,29 @@ if __name__ == "__main__":
                 log.info(f"Failed to generate gantt chart, not posting reply.{e}")
                 traceback.print_exc()
 
-        if not args.since:
+        if not since:
             log.info(
                 f"Updating last event date to {pretty_time(last_event_at)} on {LAST_MARGE_EVENT_FILE}\n"
             )
             with open(LAST_MARGE_EVENT_FILE, "w") as f:
                 f.write(last_event_at)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Monitor rejected pipelines by Marge.")
+    parser.add_argument(
+        "--token",
+        metavar="token",
+        type=str,
+        default=None,
+        help="force GitLab token, otherwise it's read from ~/.config/gitlab-token",
+    )
+    parser.add_argument(
+        "--since",
+        metavar="since",
+        type=str,
+        default=None,
+        help="consider only events after this date (ISO format), otherwise it's read from ~/.config/last_marge_event",
+    )
+    args = parser.parse_args()
+    main(args.token, args.since)