Fix tests

apache · Feb 14, 2024 · cdfa800 · cdfa800
1 parent 88ad53e
commit cdfa800
Show file tree

Hide file tree

Showing 3 changed files with 88 additions and 94 deletions.
diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py
@@ -296,7 +296,7 @@ def get_cluster(
 
     def check_cluster_autoscaling_ability(self, cluster: Cluster | dict):
         """
-        Helper method to check if the specified Cluster has ability to autoscale.
+        Check if the specified Cluster has ability to autoscale.
 
         Cluster should be Autopilot, with Node Auto-provisioning or regular auto-scaled node pools.
         Returns True if the Cluster supports autoscaling, otherwise returns False.

diff --git a/airflow/providers/google/cloud/operators/kubernetes_engine.py b/airflow/providers/google/cloud/operators/kubernetes_engine.py
@@ -76,7 +76,7 @@ def __init__(
         self._ssl_ca_cert = None
 
     def fetch_cluster_info(self) -> tuple[str, str | None]:
-        """Fetches cluster info for connecting to it."""
+        """Fetch cluster info for connecting to it."""
         cluster = self.cluster_hook.get_cluster(
             name=self.cluster_name,
             project_id=self.project_id,
@@ -530,7 +530,7 @@ def pod_hook(self) -> GKEPodHook:
 
     @staticmethod
     def _get_yaml_content_from_file(kueue_yaml_url) -> list[dict]:
-        """Helper method to download content of YAML file and separate it into several dictionaries."""
+        """Download content of YAML file and separate it into several dictionaries."""
         response = requests.get(kueue_yaml_url, allow_redirects=True)
         yaml_dicts = []
         if response.status_code == 200:
@@ -545,13 +545,12 @@ def _get_yaml_content_from_file(kueue_yaml_url) -> list[dict]:
         return yaml_dicts
 
     def execute(self, context: Context):
-        cluster_info = GKEClusterAuthDetails(
-            project_id=self.project_id,
+        self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
             cluster_name=self.cluster_name,
+            project_id=self.project_id,
             use_internal_ip=self.use_internal_ip,
             cluster_hook=self.cluster_hook,
-        )
-        self._cluster_url, self._ssl_ca_cert = cluster_info.fetch_cluster_info()
+        ).fetch_cluster_info()
 
         cluster = self.cluster_hook.get_cluster(
             name=self.cluster_name,
@@ -733,7 +732,6 @@ def hook(self) -> GKEPodHook:
         return hook
 
     def execute(self, context: Context):
-        """Executes process of creating pod and executing provided command inside it."""
         """Execute process of creating pod and executing provided command inside it."""
         self.fetch_cluster_info()
         return super().execute(context)