Used cached property for hook in SparkKubernetesOperator (#34130)

apache · Sep 6, 2023 · 1005501 · 1005501
1 parent fb99f6c
commit 1005501
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 3 deletions.
diff --git a/airflow/providers/cncf/kubernetes/operators/spark_kubernetes.py b/airflow/providers/cncf/kubernetes/operators/spark_kubernetes.py
@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import datetime
+from functools import cached_property
 from typing import TYPE_CHECKING, Sequence
 
 from kubernetes.client import ApiException
@@ -81,7 +82,9 @@ def __init__(
         self.config_file = config_file
         self.watch = watch
 
-        self.hook = KubernetesHook(
+    @cached_property
+    def hook(self) -> KubernetesHook:
+        return KubernetesHook(
             conn_id=self.kubernetes_conn_id,
             in_cluster=self.in_cluster,
             config_file=self.config_file,

diff --git a/tests/providers/cncf/kubernetes/operators/test_spark_kubernetes.py b/tests/providers/cncf/kubernetes/operators/test_spark_kubernetes.py
@@ -28,16 +28,31 @@
 
 @patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes.KubernetesHook")
 def test_spark_kubernetes_operator(mock_kubernetes_hook):
-    SparkKubernetesOperator(
+    operator = SparkKubernetesOperator(
         task_id="task_id",
         application_file="application_file",
         kubernetes_conn_id="kubernetes_conn_id",
         in_cluster=True,
         cluster_context="cluster_context",
         config_file="config_file",
     )
+    mock_kubernetes_hook.assert_not_called()  # constructor shouldn't call the hook
 
-    mock_kubernetes_hook.assert_called_once_with(
+    assert "hook" not in operator.__dict__  # Cached property has not been accessed as part of construction.
+
+
+@patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes.KubernetesHook")
+def test_spark_kubernetes_operator_hook(mock_kubernetes_hook):
+    operator = SparkKubernetesOperator(
+        task_id="task_id",
+        application_file="application_file",
+        kubernetes_conn_id="kubernetes_conn_id",
+        in_cluster=True,
+        cluster_context="cluster_context",
+        config_file="config_file",
+    )
+    operator.hook
+    mock_kubernetes_hook.assert_called_with(
         conn_id="kubernetes_conn_id",
         in_cluster=True,
         cluster_context="cluster_context",