Fix google gke system dag deprecations

providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py

@@ -447,7 +447,9 @@ def _validate_input(self) -> None:
                     "body['node_config'] or body['initial_code_count'] are specified."
                 )
 
-        if not any([self._body_field("node_config"), self._body_field("initial_node_count")]):
+        if not self._is_autopilot_cluster() and not any(
+            [self._body_field("node_config"), self._body_field("initial_node_count")]
+        ):
             if not self._body_field("node_pools"):
                 error_messages.append(
                     "Field body['node_pools'] is required if none of fields "
@@ -466,6 +468,12 @@ def _body_field(self, field_name: str, default_value: Any = None) -> Any:
             return self.body.get(field_name, default_value)
         return getattr(self.body, field_name, default_value)
 
+    def _is_autopilot_cluster(self) -> bool:
+        autopilot = self._body_field("autopilot")
+        if isinstance(autopilot, dict):
+            return bool(autopilot.get("enabled"))
+        return bool(getattr(autopilot, "enabled", False))
+
     def _alert_deprecated_body_fields(self) -> None:
         """Generate warning messages if deprecated fields were used in the body."""
         deprecated_body_fields_with_replacement = [

providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine.py

@@ -51,7 +51,7 @@
 CLUSTER_NAME = CLUSTER_NAME_BASE if len(CLUSTER_NAME_FULL) >= 33 else CLUSTER_NAME_FULL
 
 # [START howto_operator_gcp_gke_create_cluster_definition]
-CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1, "autopilot": {"enabled": True}}
+CLUSTER = {"name": CLUSTER_NAME, "autopilot": {"enabled": True}}
 # [END howto_operator_gcp_gke_create_cluster_definition]
 
 

providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py

@@ -51,7 +51,10 @@
 CLUSTER_NAME_FULL = CLUSTER_NAME_BASE + f"-{ENV_ID}".replace("_", "-")
 CLUSTER_NAME = CLUSTER_NAME_BASE if len(CLUSTER_NAME_FULL) >= 33 else CLUSTER_NAME_FULL
 
-CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1}
+CLUSTER = {
+    "name": CLUSTER_NAME,
+    "node_pools": [{"name": "default-pool", "initial_node_count": 1}],
+}
 
 with DAG(
     DAG_ID,

providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_job.py

@@ -49,7 +49,10 @@
 
 GCP_LOCATION = "europe-north1-a"
 CLUSTER_NAME = f"gke-job-{ENV_ID}".replace("_", "-")
-CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1}
+CLUSTER = {
+    "name": CLUSTER_NAME,
+    "node_pools": [{"name": "default-pool", "initial_node_count": 1}],
+}
 
 JOB_NAME = "test-pi"
 JOB_NAME_DEF = "test-pi-def"

providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_kueue.py

@@ -47,7 +47,7 @@
 
 GCP_LOCATION = "europe-west3"
 CLUSTER_NAME = f"gke-kueue-{ENV_ID}".replace("_", "-")
-CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1, "autopilot": {"enabled": True}}
+CLUSTER = {"name": CLUSTER_NAME, "autopilot": {"enabled": True}}
 
 flavor_conf = """
 apiVersion: kueue.x-k8s.io/v1beta1

providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_resource.py

@@ -44,7 +44,10 @@
 
 GCP_LOCATION = "europe-north1-a"
 CLUSTER_NAME = f"gke-resource-{ENV_ID}".replace("_", "-")
-CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1}
+CLUSTER = {
+    "name": CLUSTER_NAME,
+    "node_pools": [{"name": "default-pool", "initial_node_count": 1}],
+}
 
 PVC_NAME = "test-pvc-name"
 PVC_CONF = f"""

providers/google/tests/unit/google/cloud/operators/test_kubernetes_engine.py

@@ -102,6 +102,11 @@
 GKE_CLUSTER_CREATE_BODY_OBJECT = Cluster(
     name=GKE_CLUSTER_NAME, node_pools=[NodePool(name="a_node_pool", initial_node_count=1)]
 )
+GKE_AUTOPILOT_CLUSTER_CREATE_BODY_DICT = {
+    "name": GKE_CLUSTER_NAME,
+    "autopilot": {"enabled": True},
+}
+GKE_AUTOPILOT_CLUSTER_CREATE_BODY_OBJECT = Cluster(name=GKE_CLUSTER_NAME, autopilot={"enabled": True})
 GKE_CLUSTER_CREATE_BODY_DICT_DEPRECATED = {"name": GKE_CLUSTER_NAME, "initial_node_count": 1}
 GKE_CLUSTER_CREATE_BODY_OBJECT_DEPRECATED = Cluster(name=GKE_CLUSTER_NAME, initial_node_count=1)
 
@@ -419,7 +424,15 @@ def test_template_fields(self):
         )
         assert set(GKECreateClusterOperator.template_fields) == expected_template_fields
 
-    @pytest.mark.parametrize("body", [GKE_CLUSTER_CREATE_BODY_DICT, GKE_CLUSTER_CREATE_BODY_OBJECT])
+    @pytest.mark.parametrize(
+        "body",
+        [
+            GKE_CLUSTER_CREATE_BODY_DICT,
+            GKE_CLUSTER_CREATE_BODY_OBJECT,
+            GKE_AUTOPILOT_CLUSTER_CREATE_BODY_DICT,
+            GKE_AUTOPILOT_CLUSTER_CREATE_BODY_OBJECT,
+        ],
+    )
     def test_body(self, body):
         op = GKECreateClusterOperator(
             task_id=TEST_TASK_ID,