fixup! fixup! fixup! Bugfix: Fix overriding pod_template_file in Ku…

…bernetesExecutor
apache · Apr 5, 2021 · 855d401 · 855d401
1 parent bfc3e70
commit 855d401
Showing 1 changed file with 20 additions and 6 deletions.
diff --git a/tests/executors/test_kubernetes_executor.py b/tests/executors/test_kubernetes_executor.py
@@ -41,7 +41,7 @@
         get_base_pod_from_template,
     )
     from airflow.kubernetes import pod_generator
-    from airflow.kubernetes.pod_generator import PodGenerator
+    from airflow.kubernetes.pod_generator import PodGenerator, datetime_to_label_safe_datestring
     from airflow.utils.state import State
 except ImportError:
     AirflowKubernetesScheduler = None  # type: ignore
@@ -237,13 +237,20 @@ def test_pod_template_file_override_in_executor_config(self, mock_get_kube_clien
             assert executor.event_buffer == {}
             assert executor.task_queue.empty()
 
+            execution_date = datetime.utcnow()
+
             executor.execute_async(
-                key=('dag', 'task', datetime.utcnow(), 1),
+                key=('dag', 'task', execution_date, 1),
                 queue=None,
                 command=['airflow', 'tasks', 'run', 'true', 'some_parameter'],
                 executor_config={
                     "pod_template_file": template_file,
-                    "pod_override": k8s.V1Pod(metadata=k8s.V1ObjectMeta(labels={"release": "stable"})),
+                    "pod_override": k8s.V1Pod(
+                        metadata=k8s.V1ObjectMeta(labels={"release": "stable"}),
+                        spec=k8s.V1PodSpec(
+                            containers=[k8s.V1Container(name="base", image="airflow:3.6")],
+                        ),
+                    ),
                 },
             )
 
@@ -265,15 +272,15 @@ def test_pod_template_file_override_in_executor_config(self, mock_get_kube_clien
                         namespace="default",
                         annotations={
                             'dag_id': 'dag',
-                            'execution_date': mock.ANY,
+                            'execution_date': execution_date.isoformat(),
                             'task_id': 'task',
                             'try_number': '1',
                         },
                         labels={
                             'airflow-worker': '5',
                             'airflow_version': mock.ANY,
                             'dag_id': 'dag',
-                            'execution_date': mock.ANY,
+                            'execution_date': datetime_to_label_safe_datestring(execution_date),
                             'kubernetes_executor': 'True',
                             'mylabel': 'foo',
                             'release': 'stable',
@@ -282,7 +289,14 @@ def test_pod_template_file_override_in_executor_config(self, mock_get_kube_clien
                         },
                     ),
                     spec=k8s.V1PodSpec(
-                        containers=mock.ANY,
+                        containers=[
+                            k8s.V1Container(
+                                name="base",
+                                image="airflow:3.6",
+                                args=['airflow', 'tasks', 'run', 'true', 'some_parameter'],
+                                env=[k8s.V1EnvVar(name='AIRFLOW_IS_K8S_EXECUTOR_POD', value='True')],
+                            )
+                        ],
                         image_pull_secrets=[k8s.V1LocalObjectReference(name='airflow-registry')],
                         scheduler_name='default-scheduler',
                         security_context=k8s.V1PodSecurityContext(fs_group=50000, run_as_user=50000),