parameterize repository owner and name in WorkflowRestartChecker

Author: izaitsevfb

aws/lambda/pytorch-auto-revert/pytorch_auto_revert/workflow_checker.py

@@ -3,7 +3,6 @@
 """
 
 import logging
-import os
 from datetime import datetime, timedelta
 from typing import Dict, Set
 
@@ -13,8 +12,10 @@
 class WorkflowRestartChecker:
     """Check if workflows have been restarted using ClickHouse."""
 
-    def __init__(self):
+    def __init__(self, repo_owner: str = "pytorch", repo_name: str = "pytorch"):
         self._cache: Dict[str, bool] = {}
+        self.repo_owner = repo_owner
+        self.repo_name = repo_name
 
     def has_restarted_workflow(self, workflow_name: str, commit_sha: str) -> bool:
         """
@@ -139,9 +140,6 @@ def restart_workflow(self, workflow_name: str, commit_sha: str) -> bool:
             logging.error(f"Failed to get GitHub client: {e}")
             return False
 
-        repo_owner = os.getenv("GITHUB_REPO_OWNER", "pytorch")
-        repo_name = os.getenv("GITHUB_REPO_NAME", "pytorch")
-
         try:
             # Use trunk/{sha} tag format
             tag_ref = f"trunk/{commit_sha}"
@@ -150,15 +148,15 @@ def restart_workflow(self, workflow_name: str, commit_sha: str) -> bool:
             workflow_file_name = f"{normalized_workflow_name}.yml"
 
             # Get repo and workflow objects
-            repo = client.get_repo(f"{repo_owner}/{repo_name}")
+            repo = client.get_repo(f"{self.repo_owner}/{self.repo_name}")
             workflow = repo.get_workflow(workflow_file_name)
 
             # Dispatch the workflow
             workflow.create_dispatch(ref=tag_ref, inputs={})
 
             # Construct the workflow runs URL
             workflow_url = (
-                f"https://github.com/{repo_owner}/{repo_name}"
+                f"https://github.com/{self.repo_owner}/{self.repo_name}"
                 f"/actions/workflows/{workflow_file_name}"
                 f"?query=branch%3Atrunk%2F{commit_sha}"
             )