Repository: apache/airflow
File: providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_spark_kubernetes.py (Python)

Code before:
        mock_await_pod_completion,
        mock_fetch_requested_container_logs,
        data_file,
    ):
        """
        Verifies that find_spark_job picks a Running Spark driver pod over a non-Running pod.
        """

        task_name = "test_find_spark_job_prefers_running_pod"
        job_spec = yaml.safe_load(data_file("spark/application_template.yaml").read_text())

        mock_create_job_name.return_value = task_name
        op = SparkKubernetesOperator(
            template_spec=job_spec,
            kubernetes_conn_id="kubernetes_default_kube_config",
            task_id=task_name,
            get_logs=True,
            reattach_on_restart=True,
        )
        context = create_context(op)

        # Succeeded pod should be selected.
        succeeded_pod = mock.MagicMock()
        succeeded_pod.metadata.creation_timestamp = timezone.datetime(2025, 1, 1, tzinfo=timezone.utc)
        succeeded_pod.metadata.name = "spark-driver-running"
        succeeded_pod.metadata.labels = {"try_number": "1"}
        succeeded_pod.status.phase = "Succeeded"

        # Terminating pod should not be selected.
        terminating_pod = mock.MagicMock()
        terminating_pod.metadata.creation_timestamp = timezone.datetime(2025, 1, 1, tzinfo=timezone.utc)
        terminating_pod.metadata.deletion_timestamp = timezone.datetime(2025, 1, 2, tzinfo=timezone.utc)
        terminating_pod.metadata.name = "spark-driver-pending"
        terminating_pod.metadata.labels = {"try_number": "1"}
        terminating_pod.status.phase = "Running"

        # Failed pod should not be selected.
        failed_pod = mock.MagicMock()
        failed_pod.metadata.creation_timestamp = timezone.datetime(2025, 1, 1, tzinfo=timezone.utc)
        failed_pod.metadata.deletion_timestamp = timezone.datetime(2025, 1, 2, tzinfo=timezone.utc)
        failed_pod.metadata.name = "spark-driver-pending"
        failed_pod.metadata.labels = {"try_number": "1"}
        failed_pod.status.phase = "Failed"

        mock_get_kube_client.list_namespaced_pod.return_value.items = [
            succeeded_pod,
            terminating_pod,
            failed_pod,
        ]

        returned_pod = op.find_spark_job(context)

Review comment (suggestion):
I think you should set a `deletion_timestamp` here as well. As I believe that under your current implementation, this test will still pass if pod phase is not 'Succeeded'. `deletion_timestamp`  is the primary criterion for pod selection. Or maybe you could leave it unset. With this test, we want to assert that `find_spark_job` picks the 'Succeeded' pod even when there are multiple valid candidates i.e non-terminating pods. 

Code after:
        mock_await_pod_completion,
        mock_fetch_requested_container_logs,
        data_file,
    ):
        """
        Verifies that find_spark_job picks a Succeeded Spark driver pod over a non-Running pod.
        """

        task_name = "test_find_spark_job_prefers_succeeded_pod"
        job_spec = yaml.safe_load(data_file("spark/application_template.yaml").read_text())

        mock_create_job_name.return_value = task_name
        op = SparkKubernetesOperator(
            template_spec=job_spec,
            kubernetes_conn_id="kubernetes_default_kube_config",
            task_id=task_name,
            get_logs=True,
            reattach_on_restart=True,
        )
        context = create_context(op)

        # Succeeded pod should be selected.
        succeeded_pod = mock.MagicMock()
        succeeded_pod.metadata.creation_timestamp = timezone.datetime(2025, 1, 1, tzinfo=timezone.utc)
        succeeded_pod.metadata.name = "spark-driver"
        succeeded_pod.metadata.labels = {"try_number": "1"}
        succeeded_pod.status.phase = "Succeeded"

        # Running pod should be selected.
        running_pod = mock.MagicMock()
        running_pod.metadata.creation_timestamp = timezone.datetime(2025, 1, 1, tzinfo=timezone.utc)
        running_pod.metadata.name = "spark-driver"
        running_pod.metadata.labels = {"try_number": "1"}
        running_pod.status.phase = "Running"

        # Terminating pod should not be selected.
        terminating_pod = mock.MagicMock()
        terminating_pod.metadata.creation_timestamp = timezone.datetime(2025, 1, 1, tzinfo=timezone.utc)
        terminating_pod.metadata.deletion_timestamp = timezone.datetime(2025, 1, 2, tzinfo=timezone.utc)
        terminating_pod.metadata.name = "spark-driver"
        terminating_pod.metadata.labels = {"try_number": "1"}
        terminating_pod.status.phase = "Running"

        mock_get_kube_client.list_namespaced_pod.return_value.items = [
            terminating_pod,
            running_pod,
            succeeded_pod,
        ]

        returned_pod = op.find_spark_job(context)