test_kubernetes_executor.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

import pathlib
import random
import re
import string
import sys
from datetime import datetime, timedelta
from unittest import mock

import pytest
import yaml
from kubernetes.client import models as k8s
from kubernetes.client.rest import ApiException
from urllib3 import HTTPResponse

from airflow import AirflowException
from airflow.exceptions import PodReconciliationError
from airflow.models.taskinstancekey import TaskInstanceKey
from airflow.operators.bash import BashOperator
from airflow.operators.empty import EmptyOperator
from airflow.utils import timezone
from airflow.utils.state import State, TaskInstanceState
from tests.test_utils.config import conf_vars

try:
    from airflow.executors.kubernetes_executor import KubernetesExecutor
    from airflow.executors.kubernetes_executor_types import POD_EXECUTOR_DONE_KEY
    from airflow.executors.kubernetes_executor_utils import (
        AirflowKubernetesScheduler,
        KubernetesJobWatcher,
        ResourceVersion,
        create_pod_id,
        get_base_pod_from_template,
    )
    from airflow.kubernetes import pod_generator
    from airflow.kubernetes.kubernetes_helper_functions import (
        annotations_for_logging_task_metadata,
        annotations_to_key,
        get_logs_task_metadata,
    )
    from airflow.kubernetes.pod_generator import PodGenerator
except ImportError:
    AirflowKubernetesScheduler = None  # type: ignore


class TestAirflowKubernetesScheduler:
    @staticmethod
    def _gen_random_string(seed, str_len):
        char_list = []
        for char_seed in range(str_len):
            random.seed(str(seed) * char_seed)
            char_list.append(random.choice(string.printable))
        return "".join(char_list)

    def _cases(self):
        cases = [
            ("my_dag_id", "my-task-id"),
            ("my.dag.id", "my.task.id"),
            ("MYDAGID", "MYTASKID"),
            ("my_dag_id", "my_task_id"),
            ("mydagid" * 200, "my_task_id" * 200),
            ("my_dág_id", "my_tásk_id"),
            ("Компьютер", "niedołężność"),
            ("影師嗎", "中華民國;$"),
        ]

        cases.extend(
            [(self._gen_random_string(seed, 200), self._gen_random_string(seed, 200)) for seed in range(100)]
        )

        return cases

    @staticmethod
    def _is_valid_pod_id(name):
        regex = r"^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$"
        return len(name) <= 253 and all(ch.lower() == ch for ch in name) and re.match(regex, name)

    @staticmethod
    def _is_safe_label_value(value):
        regex = r"^[^a-z0-9A-Z]*|[^a-zA-Z0-9_\-\.]|[^a-z0-9A-Z]*$"
        return len(value) <= 63 and re.match(regex, value)

    @pytest.mark.skipif(
        AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
    )
    def test_create_pod_id(self):
        for dag_id, task_id in self._cases():
            pod_name = PodGenerator.make_unique_pod_id(create_pod_id(dag_id, task_id))
            assert self._is_valid_pod_id(pod_name)

    @pytest.mark.skipif(
        AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
    )
    @mock.patch("airflow.kubernetes.pod_generator.PodGenerator")
    @mock.patch("airflow.executors.kubernetes_executor.KubeConfig")
    def test_get_base_pod_from_template(self, mock_kubeconfig, mock_generator):
        # Provide non-existent file path,
        # so None will be passed to deserialize_model_dict().
        pod_template_file_path = "/bar/biz"
        get_base_pod_from_template(pod_template_file_path, None)
        assert "deserialize_model_dict" == mock_generator.mock_calls[0][0]
        assert mock_generator.mock_calls[0][1][0] is None

        mock_kubeconfig.pod_template_file = "/foo/bar"
        get_base_pod_from_template(None, mock_kubeconfig)
        assert "deserialize_model_dict" == mock_generator.mock_calls[1][0]
        assert mock_generator.mock_calls[1][1][0] is None

        # Provide existent file path,
        # so loaded YAML file content should be used to call deserialize_model_dict(), rather than None.
        path = sys.path[0] + "/tests/kubernetes/pod.yaml"
        with open(path) as stream:
            expected_pod_dict = yaml.safe_load(stream)

        pod_template_file_path = path
        get_base_pod_from_template(pod_template_file_path, None)
        assert "deserialize_model_dict" == mock_generator.mock_calls[2][0]
        assert mock_generator.mock_calls[2][1][0] == expected_pod_dict

        mock_kubeconfig.pod_template_file = path
        get_base_pod_from_template(None, mock_kubeconfig)
        assert "deserialize_model_dict" == mock_generator.mock_calls[3][0]
        assert mock_generator.mock_calls[3][1][0] == expected_pod_dict

    def test_make_safe_label_value(self):
        for dag_id, task_id in self._cases():
            safe_dag_id = pod_generator.make_safe_label_value(dag_id)
            assert self._is_safe_label_value(safe_dag_id)
            safe_task_id = pod_generator.make_safe_label_value(task_id)
            assert self._is_safe_label_value(safe_task_id)
            dag_id = "my_dag_id"
            assert dag_id == pod_generator.make_safe_label_value(dag_id)
            dag_id = "my_dag_id_" + "a" * 64
            assert "my_dag_id_" + "a" * 43 + "-0ce114c45" == pod_generator.make_safe_label_value(dag_id)

    def test_execution_date_serialize_deserialize(self):
        datetime_obj = datetime.now()
        serialized_datetime = pod_generator.datetime_to_label_safe_datestring(datetime_obj)
        new_datetime_obj = pod_generator.label_safe_datestring_to_datetime(serialized_datetime)

        assert datetime_obj == new_datetime_obj

    @pytest.mark.skipif(
        AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
    )
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    @mock.patch("airflow.executors.kubernetes_executor_utils.client")
    @mock.patch("airflow.executors.kubernetes_executor_utils.KubernetesJobWatcher")
    def test_delete_pod_successfully(self, mock_watcher, mock_client, mock_kube_client):
        pod_name = "my-pod-1"
        namespace = "my-namespace-1"

        mock_delete_namespace = mock.MagicMock()
        mock_kube_client.return_value.delete_namespaced_pod = mock_delete_namespace

        kube_executor = KubernetesExecutor()
        kube_executor.job_id = 1
        kube_executor.start()
        try:
            kube_executor.kube_scheduler.delete_pod(pod_name, namespace)
            mock_delete_namespace.assert_called_with(pod_name, namespace, body=mock_client.V1DeleteOptions())
        finally:
            kube_executor.end()

    @pytest.mark.skipif(
        AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
    )
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    @mock.patch("airflow.executors.kubernetes_executor_utils.client")
    @mock.patch("airflow.executors.kubernetes_executor_utils.KubernetesJobWatcher")
    def test_delete_pod_raises_404(self, mock_watcher, mock_client, mock_kube_client):
        pod_name = "my-pod-1"
        namespace = "my-namespace-2"

        mock_delete_namespace = mock.MagicMock()
        mock_kube_client.return_value.delete_namespaced_pod = mock_delete_namespace

        # ApiException is raised because status is not 404
        mock_kube_client.return_value.delete_namespaced_pod.side_effect = ApiException(status=400)
        kube_executor = KubernetesExecutor()
        kube_executor.job_id = 1
        kube_executor.start()

        with pytest.raises(ApiException):
            kube_executor.kube_scheduler.delete_pod(pod_name, namespace)
            mock_delete_namespace.assert_called_with(pod_name, namespace, body=mock_client.V1DeleteOptions())

    @pytest.mark.skipif(
        AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
    )
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    @mock.patch("airflow.executors.kubernetes_executor_utils.client")
    @mock.patch("airflow.executors.kubernetes_executor_utils.KubernetesJobWatcher")
    def test_delete_pod_404_not_raised(self, mock_watcher, mock_client, mock_kube_client):
        pod_name = "my-pod-1"
        namespace = "my-namespace-3"

        mock_delete_namespace = mock.MagicMock()
        mock_kube_client.return_value.delete_namespaced_pod = mock_delete_namespace

        # ApiException not raised because the status is 404
        mock_kube_client.return_value.delete_namespaced_pod.side_effect = ApiException(status=404)
        kube_executor = KubernetesExecutor()
        kube_executor.job_id = 1
        kube_executor.start()
        try:
            kube_executor.kube_scheduler.delete_pod(pod_name, namespace)
            mock_delete_namespace.assert_called_with(pod_name, namespace, body=mock_client.V1DeleteOptions())
        finally:
            kube_executor.end()


class TestKubernetesExecutor:
    """
    Tests if an ApiException from the Kube Client will cause the task to
    be rescheduled.
    """

    def setup_method(self) -> None:
        self.kubernetes_executor = KubernetesExecutor()
        self.kubernetes_executor.job_id = 5

    @pytest.mark.skipif(
        AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
    )
    @pytest.mark.parametrize(
        "status, should_requeue",
        [
            pytest.param(403, True, id="403 Forbidden"),
            pytest.param(12345, True, id="12345 fake-unhandled-reason"),
            pytest.param(422, False, id="422 Unprocessable Entity"),
            pytest.param(400, False, id="400 BadRequest"),
        ],
    )
    @mock.patch("airflow.executors.kubernetes_executor_utils.KubernetesJobWatcher")
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    def test_run_next_exception_requeue(
        self, mock_get_kube_client, mock_kubernetes_job_watcher, status, should_requeue
    ):
        """
        When pod scheduling fails with either reason 'Forbidden', or any reason not yet
        handled in the relevant try-except block, the task should stay in the ``task_queue``
        and be attempted on a subsequent executor sync.  When reason is 'Unprocessable Entity'
        or 'BadRequest', the task should be failed without being re-queued.

        Note on error scenarios:

        - 403 Forbidden will be returned when your request exceeds namespace quota.
        - 422 Unprocessable Entity is returned when your parameters are valid but unsupported
            e.g. limits lower than requests.
        - 400 BadRequest is returned when your parameters are invalid e.g. asking for cpu=100ABC123.

        """
        path = sys.path[0] + "/tests/kubernetes/pod_generator_base_with_secrets.yaml"

        response = HTTPResponse(body='{"message": "any message"}', status=status)

        # A mock kube_client that throws errors when making a pod
        mock_kube_client = mock.patch("kubernetes.client.CoreV1Api", autospec=True)
        mock_kube_client.create_namespaced_pod = mock.MagicMock(side_effect=ApiException(http_resp=response))
        mock_get_kube_client.return_value = mock_kube_client
        mock_api_client = mock.MagicMock()
        mock_api_client.sanitize_for_serialization.return_value = {}
        mock_kube_client.api_client = mock_api_client
        config = {
            ("kubernetes", "pod_template_file"): path,
        }
        with conf_vars(config):
            kubernetes_executor = self.kubernetes_executor
            kubernetes_executor.start()
            try:
                # Execute a task while the Api Throws errors
                try_number = 1
                task_instance_key = TaskInstanceKey("dag", "task", "run_id", try_number)
                kubernetes_executor.execute_async(
                    key=task_instance_key,
                    queue=None,
                    command=["airflow", "tasks", "run", "true", "some_parameter"],
                )
                kubernetes_executor.sync()

                assert mock_kube_client.create_namespaced_pod.call_count == 1

                if should_requeue:
                    assert not kubernetes_executor.task_queue.empty()

                    # Disable the ApiException
                    mock_kube_client.create_namespaced_pod.side_effect = None

                    # Execute the task without errors should empty the queue
                    mock_kube_client.create_namespaced_pod.reset_mock()
                    kubernetes_executor.sync()
                    assert mock_kube_client.create_namespaced_pod.called
                    assert kubernetes_executor.task_queue.empty()
                else:
                    assert kubernetes_executor.task_queue.empty()
                    assert kubernetes_executor.event_buffer[task_instance_key][0] == State.FAILED
            finally:
                kubernetes_executor.end()

    @pytest.mark.skipif(
        AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
    )
    @mock.patch("airflow.settings.pod_mutation_hook")
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    def test_run_next_pmh_error(self, mock_get_kube_client, mock_pmh):
        """
        Exception during Pod Mutation Hook execution should be handled gracefully.
        """
        exception_in_pmh = Exception("Purposely generate error for test")
        mock_pmh.side_effect = exception_in_pmh

        mock_kube_client = mock.patch("kubernetes.client.CoreV1Api", autospec=True)
        mock_kube_client.create_namespaced_pod = mock.MagicMock()
        mock_get_kube_client.return_value = mock_kube_client

        kubernetes_executor = self.kubernetes_executor
        kubernetes_executor.start()
        try:
            try_number = 1
            task_instance_key = TaskInstanceKey("dag", "task", "run_id", try_number)
            kubernetes_executor.execute_async(
                key=task_instance_key,
                queue=None,
                command=["airflow", "tasks", "run", "true", "some_parameter"],
            )
            kubernetes_executor.sync()

            # The pod_mutation_hook should have been called once.
            assert mock_pmh.call_count == 1
            # There should be no pod creation request sent
            assert mock_kube_client.create_namespaced_pod.call_count == 0
            # The task is not re-queued and there is the failed record in event_buffer
            assert kubernetes_executor.task_queue.empty()
            assert kubernetes_executor.event_buffer[task_instance_key][0] == State.FAILED
            assert kubernetes_executor.event_buffer[task_instance_key][1].__cause__ == exception_in_pmh
        finally:
            kubernetes_executor.end()

    @pytest.mark.skipif(
        AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
    )
    @mock.patch("airflow.executors.kubernetes_executor_utils.KubernetesJobWatcher")
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    def test_run_next_pod_reconciliation_error(self, mock_get_kube_client, mock_kubernetes_job_watcher):
        """
        When construct_pod raises PodReconciliationError, we should fail the task.
        """
        path = sys.path[0] + "/tests/kubernetes/pod_generator_base_with_secrets.yaml"

        mock_kube_client = mock.patch("kubernetes.client.CoreV1Api", autospec=True)
        fail_msg = "test message"
        mock_kube_client.create_namespaced_pod = mock.MagicMock(side_effect=PodReconciliationError(fail_msg))
        mock_get_kube_client.return_value = mock_kube_client
        mock_api_client = mock.MagicMock()
        mock_api_client.sanitize_for_serialization.return_value = {}
        mock_kube_client.api_client = mock_api_client
        config = {("kubernetes", "pod_template_file"): path}
        with conf_vars(config):
            kubernetes_executor = self.kubernetes_executor
            kubernetes_executor.start()
            try:
                # Execute a task while the Api Throws errors
                try_number = 1
                task_instance_key = TaskInstanceKey("dag", "task", "run_id", try_number)
                kubernetes_executor.execute_async(
                    key=task_instance_key,
                    queue=None,
                    command=["airflow", "tasks", "run", "true", "some_parameter"],
                )
                kubernetes_executor.sync()

                assert kubernetes_executor.task_queue.empty()
                assert kubernetes_executor.event_buffer[task_instance_key][0] == State.FAILED
                assert kubernetes_executor.event_buffer[task_instance_key][1].args[0] == fail_msg
            finally:
                kubernetes_executor.end()

    @mock.patch("airflow.executors.kubernetes_executor.KubeConfig")
    @mock.patch("airflow.executors.kubernetes_executor.KubernetesExecutor.sync")
    @mock.patch("airflow.executors.base_executor.BaseExecutor.trigger_tasks")
    @mock.patch("airflow.executors.base_executor.Stats.gauge")
    def test_gauge_executor_metrics(self, mock_stats_gauge, mock_trigger_tasks, mock_sync, mock_kube_config):
        executor = self.kubernetes_executor
        executor.heartbeat()
        calls = [
            mock.call(
                "executor.open_slots", value=mock.ANY, tags={"status": "open", "name": "KubernetesExecutor"}
            ),
            mock.call(
                "executor.queued_tasks",
                value=mock.ANY,
                tags={"status": "queued", "name": "KubernetesExecutor"},
            ),
            mock.call(
                "executor.running_tasks",
                value=mock.ANY,
                tags={"status": "running", "name": "KubernetesExecutor"},
            ),
        ]
        mock_stats_gauge.assert_has_calls(calls)

    @mock.patch("airflow.executors.kubernetes_executor_utils.KubernetesJobWatcher")
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    def test_invalid_executor_config(self, mock_get_kube_client, mock_kubernetes_job_watcher):
        executor = self.kubernetes_executor
        executor.start()
        try:
            assert executor.event_buffer == {}
            executor.execute_async(
                key=("dag", "task", datetime.utcnow(), 1),
                queue=None,
                command=["airflow", "tasks", "run", "true", "some_parameter"],
                executor_config=k8s.V1Pod(
                    spec=k8s.V1PodSpec(
                        containers=[k8s.V1Container(name="base", image="myimage", image_pull_policy="Always")]
                    )
                ),
            )

            assert list(executor.event_buffer.values())[0][1] == "Invalid executor_config passed"
        finally:
            executor.end()

    @pytest.mark.execution_timeout(10)
    @pytest.mark.skipif(
        AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
    )
    @mock.patch("airflow.executors.kubernetes_executor_utils.AirflowKubernetesScheduler.run_pod_async")
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    def test_pod_template_file_override_in_executor_config(self, mock_get_kube_client, mock_run_pod_async):
        current_folder = pathlib.Path(__file__).parent.resolve()
        template_file = str(
            (current_folder / "kubernetes_executor_template_files" / "basic_template.yaml").resolve()
        )

        mock_kube_client = mock.patch("kubernetes.client.CoreV1Api", autospec=True)
        mock_get_kube_client.return_value = mock_kube_client

        with conf_vars({("kubernetes", "pod_template_file"): ""}):
            executor = self.kubernetes_executor
            executor.start()
            try:
                assert executor.event_buffer == {}
                assert executor.task_queue.empty()

                executor.execute_async(
                    key=TaskInstanceKey("dag", "task", "run_id", 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"}),
                            spec=k8s.V1PodSpec(
                                containers=[k8s.V1Container(name="base", image="airflow:3.6")],
                            ),
                        ),
                    },
                )

                assert not executor.task_queue.empty()
                task = executor.task_queue.get_nowait()
                _, _, expected_executor_config, expected_pod_template_file = task
                executor.task_queue.task_done()
                # Test that the correct values have been put to queue
                assert expected_executor_config.metadata.labels == {"release": "stable"}
                assert expected_pod_template_file == template_file

                self.kubernetes_executor.kube_scheduler.run_next(task)
                mock_run_pod_async.assert_called_once_with(
                    k8s.V1Pod(
                        api_version="v1",
                        kind="Pod",
                        metadata=k8s.V1ObjectMeta(
                            name=mock.ANY,
                            namespace="default",
                            annotations={
                                "dag_id": "dag",
                                "run_id": "run_id",
                                "task_id": "task",
                                "try_number": "1",
                            },
                            labels={
                                "airflow-worker": "5",
                                "airflow_version": mock.ANY,
                                "dag_id": "dag",
                                "run_id": "run_id",
                                "kubernetes_executor": "True",
                                "mylabel": "foo",
                                "release": "stable",
                                "task_id": "task",
                                "try_number": "1",
                            },
                        ),
                        spec=k8s.V1PodSpec(
                            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),
                        ),
                    )
                )
            finally:
                executor.end()

    @mock.patch("airflow.executors.kubernetes_executor_utils.KubernetesJobWatcher")
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    def test_change_state_running(self, mock_get_kube_client, mock_kubernetes_job_watcher):
        executor = self.kubernetes_executor
        executor.start()
        try:
            key = ("dag_id", "task_id", "run_id", "try_number1")
            executor.running = {key}
            executor._change_state(key, State.RUNNING, "pod_name", "default")
            assert executor.event_buffer[key][0] == State.RUNNING
            assert executor.running == {key}
        finally:
            executor.end()

    @mock.patch("airflow.executors.kubernetes_executor_utils.KubernetesJobWatcher")
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    @mock.patch("airflow.executors.kubernetes_executor_utils.AirflowKubernetesScheduler.delete_pod")
    def test_change_state_success(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher):
        executor = self.kubernetes_executor
        executor.start()
        try:
            key = ("dag_id", "task_id", "run_id", "try_number2")
            executor.running = {key}
            executor._change_state(key, State.SUCCESS, "pod_name", "default")
            assert executor.event_buffer[key][0] == State.SUCCESS
            assert executor.running == set()
            mock_delete_pod.assert_called_once_with(pod_name="pod_name", namespace="default")
        finally:
            executor.end()

    @mock.patch("airflow.executors.kubernetes_executor_utils.KubernetesJobWatcher")
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    @mock.patch("airflow.executors.kubernetes_executor_utils.AirflowKubernetesScheduler")
    def test_change_state_failed_no_deletion(
        self, mock_kubescheduler, mock_get_kube_client, mock_kubernetes_job_watcher
    ):
        mock_delete_pod = mock_kubescheduler.return_value.delete_pod
        mock_patch_pod = mock_kubescheduler.return_value.patch_pod_executor_done
        executor = self.kubernetes_executor
        executor.kube_config.delete_worker_pods = False
        executor.kube_config.delete_worker_pods_on_failure = False
        executor.start()
        try:
            key = ("dag_id", "task_id", "run_id", "try_number3")
            executor.running = {key}
            executor._change_state(key, State.FAILED, "pod_id", "test-namespace")
            assert executor.event_buffer[key][0] == State.FAILED
            assert executor.running == set()
            mock_delete_pod.assert_not_called()
            mock_patch_pod.assert_called_once_with(pod_name="pod_id", namespace="test-namespace")
        finally:
            executor.end()

    @pytest.mark.parametrize(
        "ti_state", [TaskInstanceState.SUCCESS, TaskInstanceState.FAILED, TaskInstanceState.DEFERRED]
    )
    @mock.patch("airflow.executors.kubernetes_executor_utils.KubernetesJobWatcher")
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    @mock.patch("airflow.executors.kubernetes_executor_utils.AirflowKubernetesScheduler.delete_pod")
    def test_change_state_none(
        self,
        mock_delete_pod,
        mock_get_kube_client,
        mock_kubernetes_job_watcher,
        ti_state,
        create_task_instance,
    ):
        """Ensure that when change_state gets state=None, it looks up the TI state from the db"""
        executor = self.kubernetes_executor
        executor.start()
        try:
            ti = create_task_instance(state=ti_state)
            key = ti.key
            executor.running = {key}
            executor._change_state(key, None, "pod_name", "default")
            assert executor.event_buffer[key][0] == ti_state
            assert executor.running == set()
            mock_delete_pod.assert_called_once_with(pod_name="pod_name", namespace="default")
        finally:
            executor.end()

    @pytest.mark.parametrize(
        "multi_namespace_mode_namespace_list, watchers_keys",
        [
            pytest.param(["A", "B", "C"], ["A", "B", "C"]),
            pytest.param(None, ["ALL_NAMESPACES"]),
        ],
    )
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    def test_watchers_under_multi_namespace_mode(
        self, mock_get_kube_client, multi_namespace_mode_namespace_list, watchers_keys
    ):
        executor = self.kubernetes_executor
        executor.kube_config.multi_namespace_mode = True
        executor.kube_config.multi_namespace_mode_namespace_list = multi_namespace_mode_namespace_list
        executor.start()
        try:
            assert list(executor.kube_scheduler.kube_watchers.keys()) == watchers_keys
            assert all(
                isinstance(v, KubernetesJobWatcher) for v in executor.kube_scheduler.kube_watchers.values()
            )
        finally:
            executor.end()

    @mock.patch("airflow.executors.kubernetes_executor_utils.KubernetesJobWatcher")
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    @mock.patch("airflow.executors.kubernetes_executor_utils.AirflowKubernetesScheduler")
    def test_change_state_skip_pod_deletion(
        self, mock_kubescheduler, mock_get_kube_client, mock_kubernetes_job_watcher
    ):
        mock_delete_pod = mock_kubescheduler.return_value.delete_pod
        mock_patch_pod = mock_kubescheduler.return_value.patch_pod_executor_done
        executor = self.kubernetes_executor
        executor.kube_config.delete_worker_pods = False
        executor.kube_config.delete_worker_pods_on_failure = False

        executor.start()
        try:
            key = ("dag_id", "task_id", "run_id", "try_number2")
            executor.running = {key}
            executor._change_state(key, State.SUCCESS, "pod_name", "test-namespace")
            assert executor.event_buffer[key][0] == State.SUCCESS
            assert executor.running == set()
            mock_delete_pod.assert_not_called()
            mock_patch_pod.assert_called_once_with(pod_name="pod_name", namespace="test-namespace")
        finally:
            executor.end()

    @mock.patch("airflow.executors.kubernetes_executor_utils.KubernetesJobWatcher")
    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    @mock.patch("airflow.executors.kubernetes_executor_utils.AirflowKubernetesScheduler")
    def test_change_state_failed_pod_deletion(
        self, mock_kubescheduler, mock_get_kube_client, mock_kubernetes_job_watcher
    ):
        mock_delete_pod = mock_kubescheduler.return_value.delete_pod
        mock_patch_pod = mock_kubescheduler.return_value.patch_pod_executor_done
        executor = self.kubernetes_executor
        executor.kube_config.delete_worker_pods_on_failure = True

        executor.start()
        try:
            key = ("dag_id", "task_id", "run_id", "try_number2")
            executor.running = {key}
            executor._change_state(key, State.FAILED, "pod_name", "test-namespace")
            assert executor.event_buffer[key][0] == State.FAILED
            assert executor.running == set()
            mock_delete_pod.assert_called_once_with(pod_name="pod_name", namespace="test-namespace")
            mock_patch_pod.assert_not_called()
        finally:
            executor.end()

    @mock.patch("airflow.executors.kubernetes_executor.KubernetesExecutor.adopt_launched_task")
    @mock.patch("airflow.executors.kubernetes_executor.KubernetesExecutor._adopt_completed_pods")
    def test_try_adopt_task_instances(self, mock_adopt_completed_pods, mock_adopt_launched_task):
        executor = self.kubernetes_executor
        executor.scheduler_job_id = "10"
        ti_key = annotations_to_key(
            {
                "dag_id": "dag",
                "run_id": "run_id",
                "task_id": "task",
                "try_number": "1",
            }
        )
        mock_ti = mock.MagicMock(queued_by_job_id="1", external_executor_id="1", key=ti_key)
        pod = k8s.V1Pod(metadata=k8s.V1ObjectMeta(name="foo"))
        mock_kube_client = mock.MagicMock()
        mock_kube_client.list_namespaced_pod.return_value.items = [pod]
        executor.kube_client = mock_kube_client

        # First adoption
        reset_tis = executor.try_adopt_task_instances([mock_ti])
        mock_kube_client.list_namespaced_pod.assert_called_once_with(
            namespace="default",
            field_selector="status.phase!=Succeeded",
            label_selector="kubernetes_executor=True,airflow-worker=1,airflow_executor_done!=True",
        )
        mock_adopt_launched_task.assert_called_once_with(mock_kube_client, pod, {ti_key: mock_ti})
        mock_adopt_completed_pods.assert_called_once()
        assert reset_tis == [mock_ti]  # assume failure adopting when checking return

        # Second adoption (queued_by_job_id and external_executor_id no longer match)
        mock_kube_client.reset_mock()
        mock_adopt_launched_task.reset_mock()
        mock_adopt_completed_pods.reset_mock()

        mock_ti.queued_by_job_id = "10"  # scheduler_job would have updated this after the first adoption
        executor.scheduler_job_id = "20"
        # assume success adopting, `adopt_launched_task` pops `ti_key` from `tis_to_flush_by_key`
        mock_adopt_launched_task.side_effect = (
            lambda client, pod, tis_to_flush_by_key: tis_to_flush_by_key.pop(ti_key)
        )

        reset_tis = executor.try_adopt_task_instances([mock_ti])
        mock_kube_client.list_namespaced_pod.assert_called_once_with(
            namespace="default",
            field_selector="status.phase!=Succeeded",
            label_selector="kubernetes_executor=True,airflow-worker=10,airflow_executor_done!=True",
        )
        mock_adopt_launched_task.assert_called_once()  # Won't check args this time around as they get mutated
        mock_adopt_completed_pods.assert_called_once()
        assert reset_tis == []  # This time our return is empty - no TIs to reset

    @mock.patch("airflow.executors.kubernetes_executor.KubernetesExecutor._adopt_completed_pods")
    def test_try_adopt_task_instances_multiple_scheduler_ids(self, mock_adopt_completed_pods):
        """We try to find pods only once per scheduler id"""
        executor = self.kubernetes_executor
        mock_kube_client = mock.MagicMock()
        executor.kube_client = mock_kube_client

        mock_tis = [
            mock.MagicMock(queued_by_job_id="10", external_executor_id="1", dag_id="dag", task_id="task"),
            mock.MagicMock(queued_by_job_id="40", external_executor_id="1", dag_id="dag", task_id="task2"),
            mock.MagicMock(queued_by_job_id="40", external_executor_id="1", dag_id="dag", task_id="task3"),
        ]

        executor.try_adopt_task_instances(mock_tis)
        assert mock_kube_client.list_namespaced_pod.call_count == 2
        mock_kube_client.list_namespaced_pod.assert_has_calls(
            [
                mock.call(
                    namespace="default",
                    field_selector="status.phase!=Succeeded",
                    label_selector="kubernetes_executor=True,airflow-worker=10,airflow_executor_done!=True",
                ),
                mock.call(
                    namespace="default",
                    field_selector="status.phase!=Succeeded",
                    label_selector="kubernetes_executor=True,airflow-worker=40,airflow_executor_done!=True",
                ),
            ],
            any_order=True,
        )

    @mock.patch("airflow.executors.kubernetes_executor.KubernetesExecutor.adopt_launched_task")
    @mock.patch("airflow.executors.kubernetes_executor.KubernetesExecutor._adopt_completed_pods")
    def test_try_adopt_task_instances_no_matching_pods(
        self, mock_adopt_completed_pods, mock_adopt_launched_task
    ):
        executor = self.kubernetes_executor
        mock_ti = mock.MagicMock(queued_by_job_id="1", external_executor_id="1", dag_id="dag", task_id="task")
        mock_kube_client = mock.MagicMock()
        mock_kube_client.list_namespaced_pod.return_value.items = []
        executor.kube_client = mock_kube_client

        tis_to_flush = executor.try_adopt_task_instances([mock_ti])
        assert tis_to_flush == [mock_ti]
        mock_adopt_launched_task.assert_not_called()
        mock_adopt_completed_pods.assert_called_once()

    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    def test_adopt_launched_task(self, mock_kube_client):
        executor = self.kubernetes_executor
        executor.scheduler_job_id = "modified"
        annotations = {
            "dag_id": "dag",
            "run_id": "run_id",
            "task_id": "task",
            "try_number": "1",
        }
        ti_key = annotations_to_key(annotations)
        pod = k8s.V1Pod(
            metadata=k8s.V1ObjectMeta(name="foo", labels={"airflow-worker": "bar"}, annotations=annotations)
        )
        tis_to_flush_by_key = {ti_key: {}}

        executor.adopt_launched_task(mock_kube_client, pod=pod, tis_to_flush_by_key=tis_to_flush_by_key)
        mock_kube_client.patch_namespaced_pod.assert_called_once_with(
            body={"metadata": {"labels": {"airflow-worker": "modified"}}},
            name="foo",
            namespace=None,
        )
        assert tis_to_flush_by_key == {}
        assert executor.running == {ti_key}

    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    def test_adopt_launched_task_api_exception(self, mock_kube_client):
        """We shouldn't think we are running the task if aren't able to patch the pod"""
        executor = self.kubernetes_executor
        executor.scheduler_job_id = "modified"
        annotations = {
            "dag_id": "dag",
            "run_id": "run_id",
            "task_id": "task",
            "try_number": "1",
        }
        ti_key = annotations_to_key(annotations)
        pod = k8s.V1Pod(metadata=k8s.V1ObjectMeta(name="foo", annotations=annotations))
        tis_to_flush_by_key = {ti_key: {}}

        mock_kube_client.patch_namespaced_pod.side_effect = ApiException(status=400)
        executor.adopt_launched_task(mock_kube_client, pod=pod, tis_to_flush_by_key=tis_to_flush_by_key)
        mock_kube_client.patch_namespaced_pod.assert_called_once_with(
            body={"metadata": {"labels": {"airflow-worker": "modified"}}},
            name="foo",
            namespace=None,
        )
        assert tis_to_flush_by_key == {ti_key: {}}
        assert executor.running == set()

    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    def test_adopt_completed_pods(self, mock_kube_client):
        """We should adopt all completed pods from other schedulers"""
        executor = self.kubernetes_executor
        executor.scheduler_job_id = "modified"
        executor.kube_client = mock_kube_client
        executor.kube_config.kube_namespace = "somens"
        pod_names = ["one", "two"]

        def get_annotations(pod_name):
            return {
                "dag_id": "dag",
                "run_id": "run_id",
                "task_id": pod_name,
                "try_number": "1",
            }

        mock_kube_client.list_namespaced_pod.return_value.items = [
            k8s.V1Pod(
                metadata=k8s.V1ObjectMeta(
                    name=pod_name,
                    labels={"airflow-worker": pod_name},
                    annotations=get_annotations(pod_name),
                    namespace="somens",
                )
            )
            for pod_name in pod_names
        ]
        expected_running_ti_keys = {annotations_to_key(get_annotations(pod_name)) for pod_name in pod_names}

        executor._adopt_completed_pods(mock_kube_client)
        mock_kube_client.list_namespaced_pod.assert_called_once_with(
            namespace="somens",
            field_selector="status.phase=Succeeded",
            label_selector="kubernetes_executor=True,airflow-worker!=modified,airflow_executor_done!=True",
        )
        assert len(pod_names) == mock_kube_client.patch_namespaced_pod.call_count
        mock_kube_client.patch_namespaced_pod.assert_has_calls(
            [
                mock.call(
                    body={"metadata": {"labels": {"airflow-worker": "modified"}}},
                    name=pod_name,
                    namespace="somens",
                )
                for pod_name in pod_names
            ],
            any_order=True,
        )
        assert executor.running == expected_running_ti_keys

    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    def test_not_adopt_unassigned_task(self, mock_kube_client):
        """
        We should not adopt any tasks that were not assigned by the scheduler.
        This ensures that there is no contention over pod management.
        """

        executor = self.kubernetes_executor
        executor.scheduler_job_id = "modified"
        tis_to_flush_by_key = {"foobar": {}}
        pod = k8s.V1Pod(
            metadata=k8s.V1ObjectMeta(
                name="foo",
                labels={"airflow-worker": "bar"},
                annotations={
                    "dag_id": "dag",
                    "run_id": "run_id",
                    "task_id": "task",
                    "try_number": "1",
                },
            )
        )
        executor.adopt_launched_task(mock_kube_client, pod=pod, tis_to_flush_by_key=tis_to_flush_by_key)
        assert not mock_kube_client.patch_namespaced_pod.called
        assert tis_to_flush_by_key == {"foobar": {}}

    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    @mock.patch("airflow.executors.kubernetes_executor_utils.AirflowKubernetesScheduler.delete_pod")
    def test_cleanup_stuck_queued_tasks(self, mock_delete_pod, mock_kube_client, dag_maker, session):
        """Delete any pods associated with a task stuck in queued."""
        executor = KubernetesExecutor()
        executor.start()
        executor.scheduler_job_id = "123"
        with dag_maker(dag_id="test_cleanup_stuck_queued_tasks"):
            op = BashOperator(task_id="bash", bash_command=["echo 0", "echo 1"])
        dag_run = dag_maker.create_dagrun()
        ti = dag_run.get_task_instance(op.task_id, session)
        ti.retries = 1
        ti.state = State.QUEUED
        ti.queued_dttm = timezone.utcnow() - timedelta(minutes=30)
        ti.refresh_from_db()
        tis = [ti]
        executor.cleanup_stuck_queued_tasks(tis)
        mock_delete_pod.assert_called_once()
        assert executor.running == set()
        executor.end()

    @pytest.mark.parametrize(
        "raw_multi_namespace_mode, raw_value_namespace_list, expected_value_in_kube_config",
        [
            pytest.param("true", "A,B,C", ["A", "B", "C"]),
            pytest.param("true", "", None),
            pytest.param("false", "A,B,C", None),
            pytest.param("false", "", None),
        ],
    )
    def test_kube_config_get_namespace_list(
        self, raw_multi_namespace_mode, raw_value_namespace_list, expected_value_in_kube_config
    ):
        config = {
            ("kubernetes", "multi_namespace_mode"): raw_multi_namespace_mode,
            ("kubernetes", "multi_namespace_mode_namespace_list"): raw_value_namespace_list,
        }
        with conf_vars(config):
            executor = KubernetesExecutor()

        assert executor.kube_config.multi_namespace_mode_namespace_list == expected_value_in_kube_config

    def test_clear_not_launched_queued_tasks_not_launched(self, dag_maker, create_dummy_dag, session):
        """If a pod isn't found for a TI, reset the state to scheduled"""
        mock_kube_client = mock.MagicMock()
        mock_kube_client.list_namespaced_pod.return_value = k8s.V1PodList(items=[])

        create_dummy_dag(dag_id="test_clear", task_id="task1", with_dagrun_type=None)
        dag_run = dag_maker.create_dagrun()

        ti = dag_run.task_instances[0]
        ti.state = State.QUEUED
        ti.queued_by_job_id = 1
        session.flush()

        executor = self.kubernetes_executor
        executor.job_id = 1
        executor.kube_client = mock_kube_client
        executor.clear_not_launched_queued_tasks(session=session)

        ti.refresh_from_db()
        assert ti.state == State.SCHEDULED
        assert mock_kube_client.list_namespaced_pod.call_count == 2
        mock_kube_client.list_namespaced_pod.assert_any_call(
            namespace="default", label_selector="dag_id=test_clear,task_id=task1,airflow-worker=1,run_id=test"
        )
        # also check that we fall back to execution_date if we didn't find the pod with run_id
        execution_date_label = pod_generator.datetime_to_label_safe_datestring(ti.execution_date)
        mock_kube_client.list_namespaced_pod.assert_called_with(
            namespace="default",
            label_selector=(
                f"dag_id=test_clear,task_id=task1,airflow-worker=1,execution_date={execution_date_label}"
            ),
        )

    @pytest.mark.parametrize(
        "task_queue, kubernetes_queue",
        [
            pytest.param("default", None),
            pytest.param("kubernetes", None),
            pytest.param("kubernetes", "kubernetes"),
        ],
    )
    def test_clear_not_launched_queued_tasks_launched(
        self, dag_maker, create_dummy_dag, session, task_queue, kubernetes_queue
    ):
        """Leave the state alone if a pod already exists"""
        mock_kube_client = mock.MagicMock()
        mock_kube_client.list_namespaced_pod.return_value = k8s.V1PodList(items=["something"])

        create_dummy_dag(dag_id="test_clear", task_id="task1", with_dagrun_type=None)
        dag_run = dag_maker.create_dagrun()

        ti = dag_run.task_instances[0]
        ti.state = State.QUEUED
        ti.queued_by_job_id = 1
        ti.queue = task_queue
        session.flush()

        executor = self.kubernetes_executor
        executor.job_id = 1
        executor.kubernetes_queue = kubernetes_queue
        executor.kube_client = mock_kube_client
        executor.clear_not_launched_queued_tasks(session=session)

        ti.refresh_from_db()
        assert ti.state == State.QUEUED
        mock_kube_client.list_namespaced_pod.assert_called_once_with(
            namespace="default", label_selector="dag_id=test_clear,task_id=task1,airflow-worker=1,run_id=test"
        )

    def test_clear_not_launched_queued_tasks_mapped_task(self, dag_maker, session):
        """One mapped task has a launched pod - other does not."""

        def list_namespaced_pod(*args, **kwargs):
            if "map_index=0" in kwargs["label_selector"]:
                return k8s.V1PodList(items=["something"])
            else:
                return k8s.V1PodList(items=[])

        mock_kube_client = mock.MagicMock()
        mock_kube_client.list_namespaced_pod.side_effect = list_namespaced_pod

        with dag_maker(dag_id="test_clear"):
            op = BashOperator.partial(task_id="bash").expand(bash_command=["echo 0", "echo 1"])

        dag_run = dag_maker.create_dagrun()
        ti0 = dag_run.get_task_instance(op.task_id, session, map_index=0)
        ti0.state = State.QUEUED
        ti0.queued_by_job_id = 1

        ti1 = dag_run.get_task_instance(op.task_id, session, map_index=1)
        ti1.state = State.QUEUED
        ti1.queued_by_job_id = 1

        session.flush()

        executor = self.kubernetes_executor
        executor.job_id = 1
        executor.kube_client = mock_kube_client
        executor.clear_not_launched_queued_tasks(session=session)

        ti0.refresh_from_db()
        ti1.refresh_from_db()
        assert ti0.state == State.QUEUED
        assert ti1.state == State.SCHEDULED

        assert mock_kube_client.list_namespaced_pod.call_count == 3
        execution_date_label = pod_generator.datetime_to_label_safe_datestring(dag_run.execution_date)
        mock_kube_client.list_namespaced_pod.assert_has_calls(
            [
                mock.call(
                    namespace="default",
                    label_selector="dag_id=test_clear,task_id=bash,airflow-worker=1,map_index=0,run_id=test",
                ),
                mock.call(
                    namespace="default",
                    label_selector="dag_id=test_clear,task_id=bash,airflow-worker=1,map_index=1,run_id=test",
                ),
                mock.call(
                    namespace="default",
                    label_selector=f"dag_id=test_clear,task_id=bash,airflow-worker=1,map_index=1,"
                    f"execution_date={execution_date_label}",
                ),
            ],
            any_order=True,
        )

    def test_clear_not_launched_queued_tasks_not_launched_other_queue(
        self, dag_maker, create_dummy_dag, session
    ):
        """Queued TI has no pod, but it is not queued for the k8s executor"""
        mock_kube_client = mock.MagicMock()
        mock_kube_client.list_namespaced_pod.return_value = k8s.V1PodList(items=[])

        create_dummy_dag(dag_id="test_clear", task_id="task1", with_dagrun_type=None)
        dag_run = dag_maker.create_dagrun()

        ti = dag_run.task_instances[0]
        ti.state = State.QUEUED
        ti.queued_by_job_id = 1
        session.flush()

        executor = self.kubernetes_executor
        executor.job_id = 1
        executor.kubernetes_queue = "kubernetes"

        executor.kube_client = mock_kube_client
        executor.clear_not_launched_queued_tasks(session=session)

        ti.refresh_from_db()
        assert ti.state == State.QUEUED
        assert mock_kube_client.list_namespaced_pod.call_count == 0

    def test_clear_not_launched_queued_tasks_clear_only_by_job_id(self, dag_maker, create_dummy_dag, session):
        """clear only not launched queued  tasks which are queued by the same executor job"""
        mock_kube_client = mock.MagicMock()
        mock_kube_client.list_namespaced_pod.return_value = k8s.V1PodList(items=[])

        create_dummy_dag(dag_id="test_clear_0", task_id="task0", with_dagrun_type=None)
        dag_run = dag_maker.create_dagrun()

        ti0 = dag_run.task_instances[0]
        ti0.state = State.QUEUED
        ti0.queued_by_job_id = 1
        session.flush()

        create_dummy_dag(dag_id="test_clear_1", task_id="task1", with_dagrun_type=None)
        dag_run = dag_maker.create_dagrun()

        ti1 = dag_run.task_instances[0]
        ti1.state = State.QUEUED
        ti1.queued_by_job_id = 2
        session.flush()

        executor = self.kubernetes_executor
        executor.job_id = 1
        executor.kube_client = mock_kube_client
        executor.clear_not_launched_queued_tasks(session=session)

        ti0.refresh_from_db()
        ti1.refresh_from_db()
        assert ti0.state == State.SCHEDULED
        assert ti1.state == State.QUEUED

    @mock.patch("airflow.kubernetes.kube_client.get_kube_client")
    def test_get_task_log(self, mock_get_kube_client, create_task_instance_of_operator):
        """fetch task log from pod"""
        mock_kube_client = mock_get_kube_client.return_value

        mock_kube_client.read_namespaced_pod_log.return_value = [b"a_", b"b_", b"c_"]
        mock_pod = mock.Mock()
        mock_pod.metadata.name = "x"
        mock_kube_client.list_namespaced_pod.return_value.items = [mock_pod]
        ti = create_task_instance_of_operator(EmptyOperator, dag_id="test_k8s_log_dag", task_id="test_task")

        executor = KubernetesExecutor()
        messages, logs = executor.get_task_log(ti=ti, try_number=1)

        mock_kube_client.read_namespaced_pod_log.assert_called_once()
        assert messages == [
            "Attempting to fetch logs from pod  through kube API",
            "Found logs through kube API",
        ]
        assert logs[0] == "a_\nb_\nc_"

        mock_kube_client.reset_mock()
        mock_kube_client.read_namespaced_pod_log.side_effect = Exception("error_fetching_pod_log")

        messages, logs = executor.get_task_log(ti=ti, try_number=1)
        assert logs == [""]
        assert messages == [
            "Attempting to fetch logs from pod  through kube API",
            "Reading from k8s pod logs failed: error_fetching_pod_log",
        ]

    def test_supports_pickling(self):
        assert KubernetesExecutor.supports_pickling

    def test_supports_sentry(self):
        assert not KubernetesExecutor.supports_sentry

    def test_annotations_for_logging_task_metadata(self):
        annotations_test = {
            "dag_id": "dag",
            "run_id": "run_id",
            "task_id": "task",
            "try_number": "1",
        }
        get_logs_task_metadata.cache_clear()
        with conf_vars({("kubernetes", "logs_task_metadata"): "True"}):
            expected_annotations = {
                "dag_id": "dag",
                "run_id": "run_id",
                "task_id": "task",
                "try_number": "1",
            }
            annotations_actual = annotations_for_logging_task_metadata(annotations_test)
            assert annotations_actual == expected_annotations
        get_logs_task_metadata.cache_clear()

    def test_annotations_for_logging_task_metadata_fallback(self):
        annotations_test = {
            "dag_id": "dag",
            "run_id": "run_id",
            "task_id": "task",
            "try_number": "1",
        }
        get_logs_task_metadata.cache_clear()
        with conf_vars({("kubernetes", "logs_task_metadata"): "False"}):
            expected_annotations = "<omitted>"
            annotations_actual = annotations_for_logging_task_metadata(annotations_test)
            assert annotations_actual == expected_annotations
        get_logs_task_metadata.cache_clear()


class TestKubernetesJobWatcher:
    test_namespace = "airflow"

    def setup_method(self):
        self.watcher = KubernetesJobWatcher(
            namespace=self.test_namespace,
            watcher_queue=mock.MagicMock(),
            resource_version="0",
            scheduler_job_id="123",
            kube_config=mock.MagicMock(),
        )
        self.kube_client = mock.MagicMock()
        self.core_annotations = {
            "dag_id": "dag",
            "task_id": "task",
            "run_id": "run_id",
            "try_number": "1",
            "execution_date": None,
        }
        self.pod = k8s.V1Pod(
            metadata=k8s.V1ObjectMeta(
                name="foo",
                annotations={"airflow-worker": "bar", **self.core_annotations},
                namespace="airflow",
                resource_version="456",
                labels={},
            ),
            status=k8s.V1PodStatus(phase="Pending"),
        )
        self.events = []

    def _run(self):
        with mock.patch("airflow.executors.kubernetes_executor_utils.watch") as mock_watch:
            mock_watch.Watch.return_value.stream.return_value = self.events
            latest_resource_version = self.watcher._run(
                self.kube_client,
                self.watcher.resource_version,
                self.watcher.scheduler_job_id,
                self.watcher.kube_config,
            )
            assert self.pod.metadata.resource_version == latest_resource_version

    def assert_watcher_queue_called_once_with_state(self, state):
        self.watcher.watcher_queue.put.assert_called_once_with(
            (
                self.pod.metadata.name,
                self.watcher.namespace,
                state,
                self.core_annotations,
                self.pod.metadata.resource_version,
            )
        )

    def test_process_status_pending(self):
        self.events.append({"type": "MODIFIED", "object": self.pod})

        self._run()
        self.watcher.watcher_queue.put.assert_not_called()

    def test_process_status_pending_deleted(self):
        self.events.append({"type": "DELETED", "object": self.pod})
        self.pod.metadata.deletion_timestamp = datetime.utcnow()

        self._run()
        self.assert_watcher_queue_called_once_with_state(State.FAILED)

    def test_process_status_failed(self):
        self.pod.status.phase = "Failed"
        self.events.append({"type": "MODIFIED", "object": self.pod})

        self._run()
        self.assert_watcher_queue_called_once_with_state(State.FAILED)

    def test_process_status_succeeded(self):
        self.pod.status.phase = "Succeeded"
        self.events.append({"type": "MODIFIED", "object": self.pod})

        self._run()
        # We don't know the TI state, so we send in None
        self.assert_watcher_queue_called_once_with_state(None)

    def test_process_status_succeeded_dedup_label(self):
        self.pod.status.phase = "Succeeded"
        self.pod.metadata.labels[POD_EXECUTOR_DONE_KEY] = "True"
        self.events.append({"type": "MODIFIED", "object": self.pod})

        self._run()
        self.watcher.watcher_queue.put.assert_not_called()

    def test_process_status_succeeded_dedup_timestamp(self):
        self.pod.status.phase = "Succeeded"
        self.pod.metadata.deletion_timestamp = datetime.utcnow()
        self.events.append({"type": "MODIFIED", "object": self.pod})

        self._run()
        self.watcher.watcher_queue.put.assert_not_called()

    def test_process_status_succeeded_type_delete(self):
        self.pod.status.phase = "Succeeded"
        self.events.append({"type": "DELETED", "object": self.pod})

        self._run()
        self.watcher.watcher_queue.put.assert_not_called()

    def test_process_status_running_deleted(self):
        self.pod.status.phase = "Running"
        self.pod.metadata.deletion_timestamp = datetime.utcnow()
        self.events.append({"type": "DELETED", "object": self.pod})

        self._run()
        self.assert_watcher_queue_called_once_with_state(State.FAILED)

    def test_process_status_running(self):
        self.pod.status.phase = "Running"
        self.events.append({"type": "MODIFIED", "object": self.pod})

        self._run()
        self.watcher.watcher_queue.put.assert_not_called()

    def test_process_status_catchall(self):
        self.pod.status.phase = "Unknown"
        self.events.append({"type": "MODIFIED", "object": self.pod})

        self._run()
        self.watcher.watcher_queue.put.assert_not_called()

    @mock.patch.object(KubernetesJobWatcher, "process_error")
    def test_process_error_event_for_410(self, mock_process_error):
        message = "too old resource version: 27272 (43334)"
        self.pod.status.phase = "Pending"
        self.pod.metadata.resource_version = "0"
        mock_process_error.return_value = "0"
        raw_object = {"code": 410, "message": message}
        self.events.append({"type": "ERROR", "object": self.pod, "raw_object": raw_object})
        self._run()
        mock_process_error.assert_called_once_with(self.events[0])

    def test_process_error_event_for_raise_if_not_410(self):
        message = "Failure message"
        self.pod.status.phase = "Pending"
        raw_object = {"code": 422, "message": message, "reason": "Test"}
        self.events.append({"type": "ERROR", "object": self.pod, "raw_object": raw_object})
        error_message = (
            rf"Kubernetes failure for {raw_object['reason']} "
            rf"with code {raw_object['code']} and message: {raw_object['message']}"
        )
        with pytest.raises(AirflowException, match=error_message):
            self._run()

    def test_recover_from_resource_too_old(self):
        # too old resource
        mock_underscore_run = mock.MagicMock()

        def effect():
            yield "500"
            while True:
                yield Exception("sentinel")

        mock_underscore_run.side_effect = effect()

        self.watcher._run = mock_underscore_run

        with mock.patch("airflow.executors.kubernetes_executor_utils.get_kube_client"):
            try:
                # self.watcher._run() is mocked and return "500" as last resource_version
                self.watcher.run()
                assert False, "Should have raised Exception"
            except Exception as e:
                assert e.args == ("sentinel",)

            # both resource_version should be 0 after _run raises an exception
            assert self.watcher.resource_version == "0"
            assert ResourceVersion().resource_version[self.test_namespace] == "0"

            # check that in the next run, _run is invoked with resource_version = 0
            mock_underscore_run.reset_mock()
            try:
                self.watcher.run()
            except Exception as e:
                assert e.args == ("sentinel",)

            mock_underscore_run.assert_called_once_with(mock.ANY, "0", mock.ANY, mock.ANY)