#
# 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 contextlib
import logging
import os
import signal
import sys
from datetime import timedelta
from unittest import mock

# Leave this it is used by the test worker.
import celery.contrib.testing.tasks  # noqa: F401
import pytest
import time_machine
from celery import Celery
from celery.result import AsyncResult
from kombu.asynchronous import set_event_loop

from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.executors.base_executor import BaseExecutor
from airflow.models.dag import DAG
from airflow.models.taskinstance import TaskInstance, TaskInstanceKey
from airflow.providers.celery.executors import celery_executor, celery_executor_utils, default_celery
from airflow.providers.celery.executors.celery_executor import CeleryExecutor
from airflow.providers.common.compat.sdk import conf
from airflow.utils.state import State

from tests_common.test_utils import db
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.dag import sync_dag_to_db
from tests_common.test_utils.taskinstance import create_task_instance
from tests_common.test_utils.version_compat import (
    AIRFLOW_V_3_0_PLUS,
    AIRFLOW_V_3_1_9_PLUS,
    AIRFLOW_V_3_1_PLUS,
    AIRFLOW_V_3_2_PLUS,
    AIRFLOW_V_3_3_PLUS,
)

try:
    # Check whether a module-level function from stats is importable.
    from airflow._shared.observability.metrics.stats import gauge  # noqa: F401

    stats_reference = "airflow._shared.observability.metrics.stats"
    _executor_name_tag_key = "executor_class_name"
except ImportError:
    stats_reference = "airflow.executors.base_executor.Stats"
    _executor_name_tag_key = "name"

if AIRFLOW_V_3_0_PLUS:
    from airflow.models.dag_version import DagVersion
if AIRFLOW_V_3_2_PLUS:
    from airflow.executors.base_executor import ExecutorConf
if AIRFLOW_V_3_1_PLUS:
    from airflow.sdk import BaseOperator, timezone
else:
    from airflow.models.baseoperator import BaseOperator  # type: ignore[attr-defined,no-redef]
    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]

pytestmark = pytest.mark.db_test


@pytest.fixture(autouse=True)
def clear_cached_workload_celery_apps():
    celery_executor_utils._get_celery_app_for_workload.cache_clear()
    yield
    celery_executor_utils._get_celery_app_for_workload.cache_clear()


FAKE_EXCEPTION_MSG = "Fake Exception"


def _prepare_test_bodies():
    if "CELERY_BROKER_URLS" in os.environ:
        return [(url,) for url in os.environ["CELERY_BROKER_URLS"].split(",")]
    return [(conf.get("celery", "BROKER_URL"))]


class FakeCeleryResult:
    @property
    def state(self):
        raise Exception(FAKE_EXCEPTION_MSG)

    def task_id(self):
        return "task_id"


@contextlib.contextmanager
def _prepare_app(broker_url=None, execute=None):
    broker_url = broker_url or conf.get("celery", "BROKER_URL")

    if AIRFLOW_V_3_0_PLUS:
        execute_name = "execute_workload"
        execute = execute or celery_executor_utils.execute_workload.__wrapped__
    else:
        execute_name = "execute_command"
        execute = execute or celery_executor_utils.execute_command.__wrapped__

    test_config = dict(celery_executor_utils.get_celery_configuration())
    test_config.update({"broker_url": broker_url})
    test_app = Celery(broker_url, config_source=test_config)
    test_execute = test_app.task(execute)
    patch_app = mock.patch.object(celery_executor_utils, "app", test_app)
    patch_execute = mock.patch.object(celery_executor_utils, execute_name, test_execute)
    # Patch factory function so CeleryExecutor instances get the test app.
    patch_factory = mock.patch.object(celery_executor_utils, "create_celery_app", return_value=test_app)

    backend = test_app.backend

    if hasattr(backend, "ResultSession"):
        # Pre-create the database tables now, otherwise SQLA via Celery has a
        # race condition where it one of the subprocesses can die with "Table
        # already exists" error, because SQLA checks for which tables exist,
        # then issues a CREATE TABLE, rather than doing CREATE TABLE IF NOT
        # EXISTS.
        session = backend.ResultSession()
        session.close()

    with patch_app, patch_execute, patch_factory:
        try:
            yield test_app
        finally:
            # Clear event loop to tear down each celery instance.
            set_event_loop(None)


class TestCeleryExecutor:
    def setup_method(self) -> None:
        db.clear_db_runs()
        db.clear_db_jobs()

    def teardown_method(self) -> None:
        db.clear_db_runs()
        db.clear_db_jobs()

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="Airflow 3.2+ prefers new configuration")
    def test_sentry_integration(self):
        assert CeleryExecutor.sentry_integration == "sentry_sdk.integrations.celery.CeleryIntegration"

    @pytest.mark.skipif(AIRFLOW_V_3_2_PLUS, reason="Test only for Airflow < 3.2")
    def test_supports_sentry(self):
        assert CeleryExecutor.supports_sentry

    def test_cli_commands_vended(self):
        assert CeleryExecutor.get_cli_commands()

    def test_celery_executor_init_with_args_kwargs(self):
        """Test that CeleryExecutor properly passes args and kwargs to BaseExecutor."""
        parallelism = 50
        team_name = "test_team"

        if AIRFLOW_V_3_2_PLUS:
            # Multi-team support with ExecutorConf requires Airflow 3.2+.
            executor = celery_executor.CeleryExecutor(parallelism=parallelism, team_name=team_name)
        else:
            executor = celery_executor.CeleryExecutor(parallelism)

        assert executor.parallelism == parallelism

        if AIRFLOW_V_3_2_PLUS:
            # Multi-team support with ExecutorConf requires Airflow 3.2+.
            assert executor.team_name == team_name
            assert executor.conf.team_name == team_name

    @pytest.mark.backend("mysql", "postgres")
    def test_exception_propagation(self, caplog):
        caplog.set_level(
            logging.ERROR, logger="airflow.providers.celery.executors.celery_executor_utils.BulkStateFetcher"
        )
        with _prepare_app():
            executor = celery_executor.CeleryExecutor()
            executor.workloads = {"key": FakeCeleryResult()}
            executor.bulk_state_fetcher._get_many_using_multiprocessing(executor.workloads.values())
        assert celery_executor_utils.CELERY_FETCH_ERR_MSG_HEADER in caplog.text, caplog.record_tuples
        assert FAKE_EXCEPTION_MSG in caplog.text, caplog.record_tuples

    @mock.patch("airflow.providers.celery.executors.celery_executor.CeleryExecutor.sync")
    @mock.patch("airflow.providers.celery.executors.celery_executor.CeleryExecutor.trigger_tasks")
    @mock.patch(f"{stats_reference}.gauge")
    def test_gauge_executor_metrics(self, mock_stats_gauge, mock_trigger_tasks, mock_sync):
        executor = celery_executor.CeleryExecutor()
        executor.heartbeat()
        calls = [
            mock.call(
                "executor.open_slots",
                value=mock.ANY,
                tags={"status": "open", _executor_name_tag_key: "CeleryExecutor"},
            ),
            mock.call(
                "executor.queued_tasks",
                value=mock.ANY,
                tags={"status": "queued", _executor_name_tag_key: "CeleryExecutor"},
            ),
            mock.call(
                "executor.running_tasks",
                value=mock.ANY,
                tags={"status": "running", _executor_name_tag_key: "CeleryExecutor"},
            ),
        ]
        mock_stats_gauge.assert_has_calls(calls)

    @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Airflow 3 doesn't have execute_command anymore")
    @pytest.mark.parametrize(
        ("command", "raise_exception"),
        [
            pytest.param(["true"], True, id="wrong-command"),
            pytest.param(["airflow", "tasks"], True, id="incomplete-command"),
            pytest.param(["airflow", "tasks", "run"], False, id="complete-command"),
        ],
    )
    def test_command_validation(self, command, raise_exception):
        """Check that we validate _on the receiving_ side, not just sending side"""
        expected_context = contextlib.nullcontext()
        if raise_exception:
            expected_context = pytest.raises(
                ValueError, match=r'The command must start with \["airflow", "tasks", "run"\]\.'
            )

        with (
            mock.patch(
                "airflow.providers.celery.executors.celery_executor_utils._execute_in_subprocess"
            ) as mock_subproc,
            mock.patch(
                "airflow.providers.celery.executors.celery_executor_utils._execute_in_fork"
            ) as mock_fork,
            mock.patch("celery.app.task.Task.request") as mock_task,
        ):
            mock_task.id = "abcdef-124215-abcdef"
            with expected_context:
                celery_executor_utils.execute_command(command)
            if raise_exception:
                mock_subproc.assert_not_called()
                mock_fork.assert_not_called()
            else:
                # One of these should be called.
                assert mock_subproc.call_args == (
                    (command, "abcdef-124215-abcdef"),
                ) or mock_fork.call_args == ((command, "abcdef-124215-abcdef"),)

    @pytest.mark.backend("mysql", "postgres")
    def test_try_adopt_task_instances_none(self, clean_dags_dagruns_and_dagbundles, testing_dag_bundle):
        start_date = timezone.utcnow() - timedelta(days=2)

        with DAG("test_try_adopt_task_instances_none", schedule=None) as dag:
            task_1 = BaseOperator(task_id="task_1", start_date=start_date)

        if AIRFLOW_V_3_0_PLUS:
            sync_dag_to_db(dag)
            dag_version = DagVersion.get_latest_version(dag.dag_id)
            key1 = create_task_instance(task=task_1, run_id=None, dag_version_id=dag_version.id)
        else:
            key1 = TaskInstance(task=task_1, run_id=None)
        tis = [key1]

        executor = celery_executor.CeleryExecutor()

        assert executor.try_adopt_task_instances(tis) == tis

    @pytest.mark.backend("mysql", "postgres")
    @time_machine.travel("2020-01-01", tick=False)
    def test_try_adopt_task_instances(self, clean_dags_dagruns_and_dagbundles, testing_dag_bundle):
        start_date = timezone.utcnow() - timedelta(days=2)

        with DAG("test_try_adopt_task_instances_none", schedule=None) as dag:
            task_1 = BaseOperator(task_id="task_1", start_date=start_date)
            task_2 = BaseOperator(task_id="task_2", start_date=start_date)

        if AIRFLOW_V_3_0_PLUS:
            sync_dag_to_db(dag)
            dag_version = DagVersion.get_latest_version(dag.dag_id)
            ti1 = create_task_instance(task=task_1, run_id=None, dag_version_id=dag_version.id)
            ti2 = create_task_instance(task=task_2, run_id=None, dag_version_id=dag_version.id)
        else:
            ti1 = TaskInstance(task=task_1, run_id=None)
            ti2 = TaskInstance(task=task_2, run_id=None)
        ti1.external_executor_id = "231"
        ti1.state = State.QUEUED
        ti2.external_executor_id = "232"
        ti2.state = State.QUEUED

        tis = [ti1, ti2]

        executor = celery_executor.CeleryExecutor()
        assert executor.running == set()
        assert executor.workloads == {}

        not_adopted_tis = executor.try_adopt_task_instances(tis)

        key_1 = TaskInstanceKey(dag.dag_id, task_1.task_id, None, 0)
        key_2 = TaskInstanceKey(dag.dag_id, task_2.task_id, None, 0)
        assert executor.running == {key_1, key_2}

        assert executor.workloads == {key_1: AsyncResult("231"), key_2: AsyncResult("232")}
        assert not_adopted_tis == []

    @pytest.mark.backend("mysql", "postgres")
    @time_machine.travel("2020-01-01", tick=False)
    def test_try_adopt_with_and_without_external_executor_id(
        self, clean_dags_dagruns_and_dagbundles, testing_dag_bundle
    ):
        """TIs with an ID are adopted; TIs without are returned for reset."""
        start_date = timezone.utcnow() - timedelta(days=2)

        with DAG("test_try_adopt_mixed", schedule=None) as dag:
            task_1 = BaseOperator(task_id="task_1", start_date=start_date)
            task_2 = BaseOperator(task_id="task_2", start_date=start_date)

        if AIRFLOW_V_3_0_PLUS:
            sync_dag_to_db(dag)
            dag_version = DagVersion.get_latest_version(dag.dag_id)
            ti_with_id = create_task_instance(task=task_1, run_id=None, dag_version_id=dag_version.id)
            ti_without_id = create_task_instance(task=task_2, run_id=None, dag_version_id=dag_version.id)
        else:
            ti_with_id = TaskInstance(task=task_1, run_id=None)
            ti_without_id = TaskInstance(task=task_2, run_id=None)
        ti_with_id.external_executor_id = "231"
        ti_with_id.state = State.QUEUED
        ti_without_id.external_executor_id = None
        ti_without_id.state = State.QUEUED

        executor = celery_executor.CeleryExecutor()
        not_adopted_tis = executor.try_adopt_task_instances([ti_with_id, ti_without_id])

        key_1 = TaskInstanceKey(dag.dag_id, task_1.task_id, None, 0)
        assert key_1 in executor.running
        assert executor.workloads == {key_1: AsyncResult("231")}
        assert not_adopted_tis == [ti_without_id]

    @pytest.fixture
    def mock_celery_revoke(self):
        with _prepare_app() as app:
            app.control.revoke = mock.MagicMock()
            yield app.control.revoke

    @pytest.mark.backend("mysql", "postgres")
    @mock.patch("airflow.providers.celery.executors.celery_executor.CeleryExecutor.fail")
    def test_cleanup_stuck_queued_tasks(
        self, mock_fail, clean_dags_dagruns_and_dagbundles, testing_dag_bundle
    ):
        start_date = timezone.utcnow() - timedelta(days=2)

        with DAG("test_cleanup_stuck_queued_tasks_failed", schedule=None) as dag:
            task = BaseOperator(task_id="task_1", start_date=start_date)

        if AIRFLOW_V_3_0_PLUS:
            sync_dag_to_db(dag)
            dag_version = DagVersion.get_latest_version(task.dag.dag_id)
            ti = create_task_instance(task=task, run_id=None, dag_version_id=dag_version.id)
        else:
            ti = TaskInstance(task=task, run_id=None)
        ti.external_executor_id = "231"
        ti.state = State.QUEUED
        ti.queued_dttm = timezone.utcnow() - timedelta(minutes=30)
        ti.queued_by_job_id = 1
        tis = [ti]
        with _prepare_app() as app:
            app.control.revoke = mock.MagicMock()
            executor = celery_executor.CeleryExecutor()
            executor.job_id = 1
            executor.running = {ti.key}
            executor.workloads = {ti.key: AsyncResult("231")}
            assert executor.has_task(ti)
            with pytest.warns(AirflowProviderDeprecationWarning, match="cleanup_stuck_queued_tasks"):
                executor.cleanup_stuck_queued_tasks(tis=tis)
            executor.sync()
        assert executor.workloads == {}
        app.control.revoke.assert_called_once_with("231")
        mock_fail.assert_called()
        assert not executor.has_task(ti)

    @pytest.mark.backend("mysql", "postgres")
    @mock.patch("airflow.providers.celery.executors.celery_executor.CeleryExecutor.fail")
    def test_revoke_task(self, mock_fail, clean_dags_dagruns_and_dagbundles, testing_dag_bundle):
        start_date = timezone.utcnow() - timedelta(days=2)

        with DAG("test_revoke_task", schedule=None) as dag:
            task = BaseOperator(task_id="task_1", start_date=start_date)

        if AIRFLOW_V_3_0_PLUS:
            sync_dag_to_db(dag)
            dag_version = DagVersion.get_latest_version(task.dag.dag_id)
            ti = create_task_instance(task=task, run_id=None, dag_version_id=dag_version.id)
        else:
            ti = TaskInstance(task=task, run_id=None)
        ti.external_executor_id = "231"
        ti.state = State.QUEUED
        ti.queued_dttm = timezone.utcnow() - timedelta(minutes=30)
        ti.queued_by_job_id = 1
        tis = [ti]
        with _prepare_app() as app:
            app.control.revoke = mock.MagicMock()
            executor = celery_executor.CeleryExecutor()
            executor.job_id = 1
            executor.running = {ti.key}
            executor.workloads = {ti.key: AsyncResult("231")}
            assert executor.has_task(ti)
            for ti in tis:
                executor.revoke_task(ti=ti)
            executor.sync()
        app.control.revoke.assert_called_once_with("231")
        assert executor.workloads == {}
        assert not executor.has_task(ti)
        mock_fail.assert_not_called()

    @conf_vars({("celery", "result_backend_sqlalchemy_engine_options"): '{"pool_recycle": 1800}'})
    def test_result_backend_sqlalchemy_engine_options(self):
        import importlib

        # Scope the mock using context manager so we can clean up afterward.
        with mock.patch("celery.Celery") as mock_celery:
            # Reload celery conf to apply the new config.
            importlib.reload(default_celery)
            # reload celery_executor_utils to recreate the celery app with new config.
            importlib.reload(celery_executor_utils)

            call_args = mock_celery.call_args.kwargs.get("config_source")
            assert "database_engine_options" in call_args
            assert call_args["database_engine_options"] == {"pool_recycle": 1800}

        # Clean up: reload modules with real Celery to restore clean state for subsequent tests.
        importlib.reload(default_celery)
        importlib.reload(celery_executor_utils)


def test_operation_timeout_config():
    assert celery_executor_utils.OPERATION_TIMEOUT == 1


class MockWorkload:
    """
    A picklable object used to mock workloads sent to Celery. Can't use the mock library
    here because it's not picklable.
    """

    def apply_async(self, *args, **kwargs):
        return 1


def _exit_gracefully(signum, _):
    print(f"{os.getpid()} Exiting gracefully upon receiving signal {signum}")
    sys.exit(signum)


@pytest.fixture
def register_signals():
    """
    Register the same signals as scheduler does to test celery_executor to make sure it does not
    hang.
    """
    orig_sigint = orig_sigterm = orig_sigusr2 = signal.SIG_DFL

    orig_sigint = signal.signal(signal.SIGINT, _exit_gracefully)
    orig_sigterm = signal.signal(signal.SIGTERM, _exit_gracefully)
    orig_sigusr2 = signal.signal(signal.SIGUSR2, _exit_gracefully)

    yield

    # Restore original signal handlers after test.
    signal.signal(signal.SIGINT, orig_sigint)
    signal.signal(signal.SIGTERM, orig_sigterm)
    signal.signal(signal.SIGUSR2, orig_sigusr2)


@pytest.mark.execution_timeout(200)
@pytest.mark.quarantined
def test_send_workloads_to_celery_hang(register_signals):
    """
    Test that celery_executor does not hang after many runs.
    """
    executor = celery_executor.CeleryExecutor()

    workload = MockWorkload()
    workload_tuples_to_send = [(None, None, None, workload) for _ in range(26)]

    for _ in range(250):
        # This loop can hang on Linux if celery_executor does something wrong with
        # multiprocessing.
        results = executor._send_workloads_to_celery(workload_tuples_to_send)
        assert results == [(None, None, 1) for _ in workload_tuples_to_send]


def _has_external_executor_id_field() -> bool:
    try:
        from airflow.executors.workloads.task import TaskInstanceDTO

        return "external_executor_id" in TaskInstanceDTO.model_fields
    except (ImportError, AttributeError):
        return False


@pytest.mark.skipif(
    not _has_external_executor_id_field(),
    reason="TaskInstanceDTO.external_executor_id not available in this airflow-core version",
)
def test_send_workload_uses_external_executor_id_as_celery_task_id():
    """Pre-assigned external_executor_id is passed as task_id to apply_async()."""
    from airflow.executors.workloads.task import TaskInstanceDTO

    pre_assigned_id = "pre-assigned-uuid-for-celery"
    ti_dto = TaskInstanceDTO(
        id="00000000-0000-0000-0000-000000000001",
        dag_version_id="00000000-0000-0000-0000-000000000002",
        task_id="test_task",
        dag_id="test_dag",
        run_id="test_run",
        try_number=1,
        map_index=-1,
        pool_slots=1,
        queue="default",
        priority_weight=1,
        external_executor_id=pre_assigned_id,
    )
    key = TaskInstanceKey(
        dag_id="test_dag", task_id="test_task", run_id="test_run", map_index=-1, try_number=1
    )

    mock_result = mock.Mock(task_id=pre_assigned_id)
    mock_celery_task = mock.Mock()
    mock_celery_task.apply_async.return_value = mock_result

    mock_app = mock.Mock()
    mock_app.tasks = {"execute_workload": mock_celery_task}

    from airflow.executors.workloads import ExecuteTask

    workload = mock.Mock(spec=ExecuteTask)
    workload.ti = ti_dto
    workload.model_dump_json.return_value = "{}"

    with mock.patch(
        "airflow.providers.celery.executors.celery_executor_utils.create_celery_app",
        return_value=mock_app,
    ):
        result_key, _, result = celery_executor_utils.send_workload_to_executor(
            (key, workload, "default", None)
        )

    mock_celery_task.apply_async.assert_called_once_with(
        args=("{}",), queue="default", task_id=pre_assigned_id
    )
    assert result.task_id == pre_assigned_id


@pytest.mark.parametrize("team_name", [None, "team-a"])
def test_get_celery_app_for_workload_reuses_cache_for_same_team(team_name):
    first_app = mock.Mock()
    second_app = mock.Mock()

    with mock.patch(
        "airflow.providers.celery.executors.celery_executor_utils.create_celery_app",
        side_effect=[first_app, second_app],
    ) as mock_create_celery_app:
        assert celery_executor_utils._get_celery_app_for_workload(team_name) is first_app
        assert celery_executor_utils._get_celery_app_for_workload(team_name) is first_app

    mock_create_celery_app.assert_called_once()


def test_get_celery_app_for_workload_keeps_cache_team_scoped():
    team_a_app = mock.Mock()
    team_b_app = mock.Mock()

    with mock.patch(
        "airflow.providers.celery.executors.celery_executor_utils.create_celery_app",
        side_effect=[team_a_app, team_b_app],
    ) as mock_create_celery_app:
        assert celery_executor_utils._get_celery_app_for_workload("team-a") is team_a_app
        assert celery_executor_utils._get_celery_app_for_workload("team-b") is team_b_app

    assert mock_create_celery_app.call_count == 2


def test_send_workload_reuses_celery_app_for_same_team():
    """Publishing multiple workloads for the same team reuses the cached Celery app."""
    key = TaskInstanceKey(
        dag_id="test_dag", task_id="test_task", run_id="test_run", map_index=-1, try_number=1
    )
    mock_result = mock.Mock(task_id="mock-task-id")
    mock_celery_task = mock.Mock()
    mock_celery_task.apply_async.return_value = mock_result
    mock_app = mock.Mock()
    task_name = "execute_workload" if AIRFLOW_V_3_0_PLUS else "execute_command"
    mock_app.tasks = {task_name: mock_celery_task}

    if AIRFLOW_V_3_0_PLUS:
        workload = mock.Mock()
        workload.ti.external_executor_id = None
        workload.model_dump_json.return_value = "{}"
    else:
        workload = ["airflow", "tasks", "run", "test_dag", "test_task", "test_run"]

    with mock.patch(
        "airflow.providers.celery.executors.celery_executor_utils.create_celery_app",
        return_value=mock_app,
    ) as mock_create_celery_app:
        celery_executor_utils.send_workload_to_executor((key, workload, "default", "team-a"))
        celery_executor_utils.send_workload_to_executor((key, workload, "default", "team-a"))

    mock_create_celery_app.assert_called_once()
    assert mock_celery_task.apply_async.call_count == 2


def test_send_workload_keeps_celery_app_cache_team_scoped():
    """Different teams get distinct cached Celery app instances in the publisher process."""
    key = TaskInstanceKey(
        dag_id="test_dag", task_id="test_task", run_id="test_run", map_index=-1, try_number=1
    )
    mock_result = mock.Mock(task_id="mock-task-id")
    team_a_task = mock.Mock()
    team_a_task.apply_async.return_value = mock_result
    team_b_task = mock.Mock()
    team_b_task.apply_async.return_value = mock_result
    task_name = "execute_workload" if AIRFLOW_V_3_0_PLUS else "execute_command"
    team_a_app = mock.Mock(tasks={task_name: team_a_task})
    team_b_app = mock.Mock(tasks={task_name: team_b_task})

    if AIRFLOW_V_3_0_PLUS:
        workload = mock.Mock()
        workload.ti.external_executor_id = None
        workload.model_dump_json.return_value = "{}"
    else:
        workload = ["airflow", "tasks", "run", "test_dag", "test_task", "test_run"]

    with mock.patch(
        "airflow.providers.celery.executors.celery_executor_utils.create_celery_app",
        side_effect=[team_a_app, team_b_app],
    ) as mock_create_celery_app:
        celery_executor_utils.send_workload_to_executor((key, workload, "default", "team-a"))
        celery_executor_utils.send_workload_to_executor((key, workload, "default", "team-b"))

    assert mock_create_celery_app.call_count == 2
    team_a_task.apply_async.assert_called_once()
    team_b_task.apply_async.assert_called_once()


@conf_vars({("celery", "result_backend"): "rediss://test_user:test_password@localhost:6379/0"})
def test_celery_executor_with_no_recommended_result_backend(caplog):
    import importlib

    from airflow.providers.celery.executors.default_celery import log

    with caplog.at_level(logging.WARNING, logger=log.name):
        # Reload celery conf to apply the new config.
        importlib.reload(default_celery)
        assert "test_password" not in caplog.text
        assert (
            "You have configured a result_backend using the protocol `rediss`,"
            " it is highly recommended to use an alternative result_backend (i.e. a database)."
        ) in caplog.text


@conf_vars({("celery_broker_transport_options", "sentinel_kwargs"): '{"service_name": "mymaster"}'})
def test_sentinel_kwargs_loaded_from_string():
    import importlib

    # Reload celery conf to apply the new config.
    importlib.reload(default_celery)
    assert default_celery.DEFAULT_CELERY_CONFIG["broker_transport_options"]["sentinel_kwargs"] == {
        "service_name": "mymaster"
    }


@conf_vars({("celery", "task_acks_late"): "False"})
def test_celery_task_acks_late_loaded_from_string():
    import importlib

    # Reload celery conf to apply the new config.
    importlib.reload(default_celery)
    assert default_celery.DEFAULT_CELERY_CONFIG["task_acks_late"] is False


@conf_vars({("celery", "BROKER_URL"): "redis://localhost:6379/0"})
def test_visibility_timeout_default_warns_when_not_configured(caplog):
    """Test that a warning is logged when visibility_timeout defaults to 86400 (24h)."""
    import importlib

    from airflow.providers.celery.executors.default_celery import log

    with caplog.at_level(logging.WARNING, logger=log.name):
        importlib.reload(default_celery)
        assert default_celery.DEFAULT_CELERY_CONFIG["broker_transport_options"]["visibility_timeout"] == 86400
        assert "No visibility_timeout configured" in caplog.text
        assert "86400" in caplog.text
        assert "long-running tasks" in caplog.text


@conf_vars(
    {
        ("celery", "BROKER_URL"): "redis://localhost:6379/0",
        ("celery_broker_transport_options", "visibility_timeout"): "172800",
    }
)
def test_visibility_timeout_no_warning_when_configured(caplog):
    """Test that no warning is logged when visibility_timeout is explicitly configured."""
    import importlib

    from airflow.providers.celery.executors.default_celery import log

    with caplog.at_level(logging.WARNING, logger=log.name):
        importlib.reload(default_celery)
        assert (
            int(default_celery.DEFAULT_CELERY_CONFIG["broker_transport_options"]["visibility_timeout"])
            == 172800
        )
        assert "No visibility_timeout configured" not in caplog.text


@conf_vars({("celery", "BROKER_URL"): "amqp://guest:guest@localhost:5672//"})
def test_visibility_timeout_not_set_for_unsupported_broker(caplog):
    """Test that visibility_timeout is not set for brokers that don't support it (e.g. RabbitMQ)."""
    import importlib

    from airflow.providers.celery.executors.default_celery import log

    with caplog.at_level(logging.WARNING, logger=log.name):
        importlib.reload(default_celery)
        assert "visibility_timeout" not in default_celery.DEFAULT_CELERY_CONFIG.get(
            "broker_transport_options", {}
        )
        assert "No visibility_timeout configured" not in caplog.text


@conf_vars({("celery", "extra_celery_config"): '{"worker_max_tasks_per_child": 10}'})
def test_celery_extra_celery_config_loaded_from_string():
    import importlib

    # Reload celery conf to apply the new config.
    importlib.reload(default_celery)
    assert default_celery.DEFAULT_CELERY_CONFIG["worker_max_tasks_per_child"] == 10


@conf_vars({("celery_result_backend_transport_options", "sentinel_kwargs"): '{"password": "redis_password"}'})
def test_result_backend_sentinel_kwargs_loaded_from_string():
    """Test that sentinel_kwargs for result backend transport options is correctly parsed."""
    import importlib

    # Reload celery conf to apply the new config.
    importlib.reload(default_celery)
    assert "result_backend_transport_options" in default_celery.DEFAULT_CELERY_CONFIG
    assert default_celery.DEFAULT_CELERY_CONFIG["result_backend_transport_options"]["sentinel_kwargs"] == {
        "password": "redis_password"
    }


@conf_vars({("celery_result_backend_transport_options", "master_name"): "mymaster"})
def test_result_backend_master_name_loaded():
    """Test that master_name for result backend transport options is correctly loaded."""
    import importlib

    # Reload celery conf to apply the new config.
    importlib.reload(default_celery)
    assert "result_backend_transport_options" in default_celery.DEFAULT_CELERY_CONFIG
    assert (
        default_celery.DEFAULT_CELERY_CONFIG["result_backend_transport_options"]["master_name"] == "mymaster"
    )


@conf_vars(
    {
        ("celery_result_backend_transport_options", "sentinel_kwargs"): '{"password": "redis_password"}',
        ("celery_result_backend_transport_options", "master_name"): "mymaster",
    }
)
def test_result_backend_transport_options_with_multiple_options():
    """Test that multiple result backend transport options are correctly loaded."""
    import importlib

    # Reload celery conf to apply the new config.
    importlib.reload(default_celery)
    result_backend_opts = default_celery.DEFAULT_CELERY_CONFIG["result_backend_transport_options"]
    assert result_backend_opts["sentinel_kwargs"] == {"password": "redis_password"}
    assert result_backend_opts["master_name"] == "mymaster"


@conf_vars({("celery_result_backend_transport_options", "sentinel_kwargs"): "invalid_json"})
def test_result_backend_sentinel_kwargs_invalid_json():
    """Test that invalid JSON in sentinel_kwargs raises an error."""
    import importlib

    from airflow.providers.common.compat.sdk import AirflowException

    with pytest.raises(
        AirflowException, match="sentinel_kwargs.*should be written in the correct dictionary format"
    ):
        importlib.reload(default_celery)


@conf_vars({("celery_result_backend_transport_options", "sentinel_kwargs"): '"not_a_dict"'})
def test_result_backend_sentinel_kwargs_not_dict():
    """Test that non-dict sentinel_kwargs raises an error."""
    import importlib

    from airflow.providers.common.compat.sdk import AirflowException

    with pytest.raises(
        AirflowException, match="sentinel_kwargs.*should be written in the correct dictionary format"
    ):
        importlib.reload(default_celery)


@conf_vars(
    {
        ("celery", "result_backend"): "sentinel://sentinel1:26379;sentinel://sentinel2:26379",
        ("celery_result_backend_transport_options", "sentinel_kwargs"): '{"password": "redis_pass"}',
        ("celery_result_backend_transport_options", "master_name"): "mymaster",
    }
)
def test_result_backend_sentinel_full_config():
    """Test full Redis Sentinel configuration for result backend."""
    import importlib

    # Reload celery conf to apply the new config.
    importlib.reload(default_celery)

    assert default_celery.DEFAULT_CELERY_CONFIG["result_backend"] == (
        "sentinel://sentinel1:26379;sentinel://sentinel2:26379"
    )
    result_backend_opts = default_celery.DEFAULT_CELERY_CONFIG["result_backend_transport_options"]
    assert result_backend_opts["sentinel_kwargs"] == {"password": "redis_pass"}
    assert result_backend_opts["master_name"] == "mymaster"


@pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="Multi-team support requires Airflow 3.2+")
class TestMultiTeamCeleryExecutor:
    """Test multi-team functionality in CeleryExecutor."""

    def setup_method(self) -> None:
        db.clear_db_runs()
        db.clear_db_jobs()

    def teardown_method(self) -> None:
        db.clear_db_runs()
        db.clear_db_jobs()

    @conf_vars(
        {
            ("celery", "broker_url"): "redis://global:6379/0",
            ("operators", "default_queue"): "global_queue",
        }
    )
    def test_multi_team_isolation_and_workload_routing(self, monkeypatch):
        """
        Test multi-team executor isolation and correct workload routing.

        Verifies:
        - Each executor has isolated Celery app and config
        - Workloads are routed through team-specific apps (_process_tasks/_process_workloads)
        - Backward compatibility with global executor
        """
        # Set up team-specific config via environment variables
        monkeypatch.setenv("AIRFLOW__TEAM_A___CELERY__BROKER_URL", "redis://team-a:6379/0")
        monkeypatch.setenv("AIRFLOW__TEAM_A___OPERATORS__DEFAULT_QUEUE", "team_a_queue")
        monkeypatch.setenv("AIRFLOW__TEAM_B___CELERY__BROKER_URL", "redis://team-b:6379/0")
        monkeypatch.setenv("AIRFLOW__TEAM_B___OPERATORS__DEFAULT_QUEUE", "team_b_queue")

        # Reload config to pick up environment variables.
        from airflow import configuration

        configuration.conf.read_dict({}, source="test")

        # Create executors with different team configs.
        team_a_executor = CeleryExecutor(parallelism=2, team_name="team_a")
        team_b_executor = CeleryExecutor(parallelism=3, team_name="team_b")
        global_executor = CeleryExecutor(parallelism=4)

        # Each executor has its own Celery app (critical for isolation).
        assert team_a_executor.celery_app is not team_b_executor.celery_app
        assert team_a_executor.celery_app is not global_executor.celery_app

        # Team-specific broker URLs are used.
        assert "team-a" in team_a_executor.celery_app.conf.broker_url
        assert "team-b" in team_b_executor.celery_app.conf.broker_url
        assert "global" in global_executor.celery_app.conf.broker_url

        # Team-specific queues are used.
        assert team_a_executor.celery_app.conf.task_default_queue == "team_a_queue"
        assert team_b_executor.celery_app.conf.task_default_queue == "team_b_queue"
        assert global_executor.celery_app.conf.task_default_queue == "global_queue"

        # Each executor has its own BulkStateFetcher with correct app.
        assert team_a_executor.bulk_state_fetcher.celery_app is team_a_executor.celery_app
        assert team_b_executor.bulk_state_fetcher.celery_app is team_b_executor.celery_app

        # Executors have isolated internal state.
        assert team_a_executor.workloads is not team_b_executor.workloads
        assert team_a_executor.running is not team_b_executor.running
        assert team_a_executor.queued_tasks is not team_b_executor.queued_tasks

    @conf_vars({("celery", "broker_url"): "redis://global:6379/0"})
    @mock.patch("airflow.providers.celery.executors.celery_executor.CeleryExecutor._send_workloads")
    def test_workload_routing_through_team_specific_app(self, mock_send_workloads, monkeypatch):
        """
        Test that _process_tasks (v2) and _process_workloads (v3) pass the correct team_name for task routing.

        With the ProcessPoolExecutor approach, we pass team_name instead of task objects to avoid
        pickling issues. The subprocess reconstructs the team-specific Celery app from the team_name.
        """
        # Set up team A config.
        monkeypatch.setenv("AIRFLOW__TEAM_A___CELERY__BROKER_URL", "redis://team-a:6379/0")

        team_a_executor = CeleryExecutor(parallelism=2, team_name="team_a")

        if AIRFLOW_V_3_0_PLUS:
            from airflow.executors.workloads import ExecuteTask
            from airflow.models.taskinstancekey import TaskInstanceKey

            # Create mock workload.
            mock_ti = mock.Mock()
            mock_ti.key = TaskInstanceKey("dag", "task", "run", 1)
            mock_ti.queue = "test_queue"
            mock_workload = mock.Mock(spec=ExecuteTask)
            mock_workload.ti = mock_ti

            # Process workload through team A executor.
            team_a_executor._process_workloads([mock_workload])

            # Verify _send_workloads received the correct team_name.
            assert mock_send_workloads.called
            workload_tuples = mock_send_workloads.call_args[0][0]
            team_name_from_call = workload_tuples[0][
                3
            ]  # 4th element is team_name (used to reconstruct Celery app in subprocess).

            # Critical: team_name is passed so subprocess can reconstruct the correct app.
            assert team_name_from_call == "team_a"
        else:
            from airflow.models.taskinstancekey import TaskInstanceKey

            # Test V2 path with execute_command.
            mock_key = TaskInstanceKey("dag", "task", "run", 1)
            mock_command = ["airflow", "tasks", "run", "dag", "task"]
            mock_queue = "test_queue"

            # Process task through team A executor.
            team_a_executor._process_tasks([(mock_key, mock_command, mock_queue, None)])

            # Verify _send_workloads received team A's execute_command workload (v2 compatibility path).
            assert mock_send_workloads.called
            task_tuples = mock_send_workloads.call_args[0][0]
            task_from_call = task_tuples[0][3]  # 4th element is the task (V2 still uses task object).

            # Critical: Celery task belongs to team A's app, not module-level app.
            assert task_from_call.app is team_a_executor.celery_app
            assert task_from_call.name == "execute_command"


def test_celery_tasks_registered_on_import():
    """
    Ensure execute_workload (and execute_command for Airflow 2.x) are registered
    with the Celery app when celery_executor is imported.

    Regression test for https://github.com/apache/airflow/issues/63043
    Celery provider 3.17.0 exposed that celery_executor_utils was never imported
    at module level, so tasks were never registered at worker startup.
    """
    from airflow.providers.celery.executors.celery_executor_utils import app

    registered_tasks = list(app.tasks.keys())
    assert "execute_workload" in registered_tasks, (
        "execute_workload must be registered with the Celery app at import time. "
        "Workers need this to receive tasks without KeyError."
    )
    # TODO: remove this block when min supported Airflow version is >= 3.0.
    if not AIRFLOW_V_3_0_PLUS:
        assert "execute_command" in registered_tasks, (
            "execute_command must be registered for Airflow 2.x compatibility."
        )


@pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="ExecuteCallback requires Airflow 3.2+")
@pytest.mark.parametrize(
    ("callback_data", "expected_queue"),
    [
        pytest.param({"path": "tests.callbacks.test_callback", "kwargs": {}}, "default", id="default_queue"),
        pytest.param(
            {"path": "tests.callbacks.test_callback", "kwargs": {}, "queue": "callback_queue"},
            "callback_queue",
            id="callback_queue",
        ),
    ],
)
@mock.patch("airflow.providers.celery.executors.celery_executor.CeleryExecutor._send_workloads")
def test_process_workloads_routes_execute_callback(mock_send_workloads, callback_data, expected_queue):
    """CeleryExecutor routes callback workloads to Celery with the expected queue."""
    from airflow.executors import workloads
    from airflow.executors.workloads.callback import CallbackDTO

    callback_id = "00000000-0000-0000-0000-000000000003"
    workload = workloads.ExecuteCallback(
        callback=CallbackDTO(
            id=callback_id,
            fetch_method=workloads.CallbackFetchMethod.IMPORT_PATH,
            data=callback_data,
        ),
        dag_rel_path="test_dag.py",
        bundle_info=workloads.BundleInfo(name="test-bundle", version=None),
        token="test-token",
        log_path="callback.log",
    )

    executor = celery_executor.CeleryExecutor()
    executor._process_workloads([workload])

    mock_send_workloads.assert_called_once_with([(workload.callback.key, workload, expected_queue, None)])


@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="execute_workload is only used for Airflow 3+")
@pytest.mark.skipif(AIRFLOW_V_3_3_PLUS, reason="pre-3.3 compatibility path only applies before Airflow 3.3")
def test_execute_workload_runs_execute_task_before_airflow_3_3():
    """execute_workload routes serialized ExecuteTask payloads to supervise on Airflow 3.0-3.2."""
    from airflow.executors import workloads

    workload = workloads.ExecuteTask(
        ti=workloads.TaskInstance(
            id="00000000-0000-0000-0000-000000000001",
            dag_version_id="00000000-0000-0000-0000-000000000002",
            task_id="test_task",
            dag_id="test_dag",
            run_id="test_run",
            try_number=1,
            map_index=-1,
            pool_slots=1,
            queue="default",
            priority_weight=1,
        ),
        dag_rel_path="test_dag.py",
        bundle_info=workloads.BundleInfo(name="test-bundle", version=None),
        token="test-token",
        log_path="test.log",
    )
    mock_current_task = mock.MagicMock()
    mock_current_task.request.id = "test-celery-task-id"
    mock_app = mock.MagicMock()
    mock_app.current_task = mock_current_task

    with (
        mock.patch.object(celery_executor_utils, "app", mock_app),
        mock.patch("airflow.sdk.execution_time.supervisor.supervise") as mock_supervise,
    ):
        celery_executor_utils.execute_workload.__wrapped__(workload.model_dump_json())

    mock_supervise.assert_called_once()
    assert mock_supervise.call_args.kwargs["ti"].task_id == "test_task"
    assert str(mock_supervise.call_args.kwargs["dag_rel_path"]) == "test_dag.py"
    assert mock_supervise.call_args.kwargs["token"] == "test-token"
    assert mock_supervise.call_args.kwargs["log_path"] == "test.log"


@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="3.3+ path uses BaseExecutor.run_workload")
def test_execute_workload_runs_base_executor_workload_on_airflow_3_3_plus():
    """execute_workload routes serialized ExecuteTask payloads to BaseExecutor on Airflow 3.3+."""
    from airflow.executors import workloads

    workload = workloads.ExecuteTask(
        ti=workloads.TaskInstance(
            id="00000000-0000-0000-0000-000000000001",
            dag_version_id="00000000-0000-0000-0000-000000000002",
            task_id="test_task",
            dag_id="test_dag",
            run_id="test_run",
            try_number=1,
            map_index=-1,
            pool_slots=1,
            queue="default",
            priority_weight=1,
        ),
        dag_rel_path="test_dag.py",
        bundle_info=workloads.BundleInfo(name="test-bundle", version=None),
        token="test-token",
        log_path="test.log",
    )
    mock_current_task = mock.MagicMock()
    mock_current_task.request.id = "test-celery-task-id"
    mock_app = mock.MagicMock()
    mock_app.current_task = mock_current_task

    with (
        mock.patch.object(celery_executor_utils, "app", mock_app),
        mock.patch("airflow.executors.base_executor.BaseExecutor.run_workload") as mock_run_workload,
    ):
        celery_executor_utils.execute_workload.__wrapped__(workload.model_dump_json())

    mock_run_workload.assert_called_once()
    decoded_workload = mock_run_workload.call_args.args[0]
    assert decoded_workload.ti.task_id == "test_task"
    assert str(decoded_workload.dag_rel_path) == "test_dag.py"
    assert decoded_workload.token == "test-token"
    assert decoded_workload.log_path == "test.log"


@pytest.mark.skipif(not AIRFLOW_V_3_1_9_PLUS, reason="TaskAlreadyRunningError requires Airflow 3.1.9+")
@pytest.mark.skipif(
    not hasattr(BaseExecutor, "run_workload"),
    reason="BaseExecutor.run_workload not available in this Airflow version",
)
def test_execute_workload_ignores_already_running_task():
    """Test that execute_workload raises Celery Ignore when task is already running."""
    import importlib

    from celery.exceptions import Ignore

    from airflow.sdk.exceptions import TaskAlreadyRunningError

    importlib.reload(celery_executor_utils)
    execute_workload_unwrapped = celery_executor_utils.execute_workload.__wrapped__

    mock_current_task = mock.MagicMock()
    mock_current_task.request.id = "test-celery-task-id"
    mock_app = mock.MagicMock()
    mock_app.current_task = mock_current_task

    with (
        mock.patch("airflow.executors.base_executor.BaseExecutor.run_workload") as mock_run_workload,
        mock.patch.object(celery_executor_utils, "app", mock_app),
    ):
        mock_run_workload.side_effect = TaskAlreadyRunningError("Task already running")

        workload_json = """
        {
            "type": "ExecuteTask",
            "token": "test-token",
            "dag_rel_path": "test_dag.py",
            "bundle_info": {"name": "test-bundle", "version": null},
            "log_path": "test.log",
            "ti": {
                "id": "019bdec0-d353-7b68-abe0-5ac20fa75ad0",
                "dag_version_id": "019bdead-fdcd-78ab-a9f2-aba3b80fded2",
                "task_id": "test_task",
                "dag_id": "test_dag",
                "run_id": "test_run",
                "try_number": 1,
                "map_index": -1,
                "pool_slots": 1,
                "queue": "default",
                "priority_weight": 1
            }
        }
        """
        with pytest.raises(Ignore):
            execute_workload_unwrapped(workload_json)


class TestAmqpsSslConfig:
    """Tests for amqps:// broker URL SSL configuration (Fix for substring match bug)."""

    @conf_vars(
        {
            ("celery", "BROKER_URL"): "amqps://guest:guest@rabbitmq:5671//",
            ("celery", "SSL_ACTIVE"): "True",
            ("celery", "SSL_KEY"): "/path/to/key.pem",
            ("celery", "SSL_CERT"): "/path/to/cert.pem",
            ("celery", "SSL_CACERT"): "/path/to/ca.pem",
        }
    )
    def test_amqps_broker_url_builds_ssl_config(self):
        """Test that amqps:// broker URLs correctly build broker_use_ssl with AMQP param names."""
        import importlib
        import ssl

        importlib.reload(default_celery)

        config = default_celery.DEFAULT_CELERY_CONFIG
        assert "broker_use_ssl" in config, "broker_use_ssl should be set for amqps:// URLs"
        broker_ssl = config["broker_use_ssl"]
        assert broker_ssl["keyfile"] == "/path/to/key.pem"
        assert broker_ssl["certfile"] == "/path/to/cert.pem"
        assert broker_ssl["ca_certs"] == "/path/to/ca.pem"
        assert broker_ssl["cert_reqs"] == ssl.CERT_REQUIRED
        # Must NOT have ssl_ prefixed keys (those are for Redis)
        assert "ssl_keyfile" not in broker_ssl
        assert "ssl_certfile" not in broker_ssl

    @conf_vars(
        {
            ("celery", "BROKER_URL"): "amqp://guest:guest@rabbitmq:5672//",
            ("celery", "SSL_ACTIVE"): "True",
            ("celery", "SSL_KEY"): "/path/to/key.pem",
            ("celery", "SSL_CERT"): "/path/to/cert.pem",
            ("celery", "SSL_CACERT"): "/path/to/ca.pem",
        }
    )
    def test_amqp_broker_url_still_builds_ssl_config(self):
        """Test that amqp:// (non-TLS) broker URLs still build SSL config correctly (no regression)."""
        import importlib
        import ssl

        importlib.reload(default_celery)

        config = default_celery.DEFAULT_CELERY_CONFIG
        assert "broker_use_ssl" in config
        broker_ssl = config["broker_use_ssl"]
        assert broker_ssl["keyfile"] == "/path/to/key.pem"
        assert broker_ssl["cert_reqs"] == ssl.CERT_REQUIRED

    @conf_vars(
        {
            ("celery", "BROKER_URL"): "rediss://redis:6380//",
            ("celery", "SSL_ACTIVE"): "True",
            ("celery", "SSL_KEY"): "/path/to/key.pem",
            ("celery", "SSL_CERT"): "/path/to/cert.pem",
            ("celery", "SSL_CACERT"): "/path/to/ca.pem",
        }
    )
    def test_redis_mutual_tls_builds_ssl_config(self):
        """Test mutual TLS: all three SSL keys produce correct broker_use_ssl for Redis."""
        import importlib
        import ssl

        importlib.reload(default_celery)

        config = default_celery.DEFAULT_CELERY_CONFIG
        assert "broker_use_ssl" in config
        broker_ssl = config["broker_use_ssl"]
        assert broker_ssl["ssl_keyfile"] == "/path/to/key.pem"
        assert broker_ssl["ssl_certfile"] == "/path/to/cert.pem"
        assert broker_ssl["ssl_ca_certs"] == "/path/to/ca.pem"
        assert broker_ssl["ssl_cert_reqs"] == ssl.CERT_REQUIRED

    @conf_vars(
        {
            ("celery", "BROKER_URL"): "amqps://guest:guest@rabbitmq:5671//",
            ("celery", "SSL_ACTIVE"): "True",
            ("celery", "SSL_CACERT"): "/path/to/ca.pem",
        }
    )
    def test_amqps_mutual_tls_missing_key_cert_raises(self):
        """Test that mutual TLS (default) raises error when SSL_KEY/SSL_CERT are missing."""
        import importlib

        with pytest.raises(ValueError, match="SSL_MUTUAL_TLS is True.*but SSL_KEY and/or SSL_CERT"):
            importlib.reload(default_celery)

    @conf_vars(
        {
            ("celery", "BROKER_URL"): "amqps://guest:guest@rabbitmq:5671//",
            ("celery", "SSL_ACTIVE"): "True",
            ("celery", "SSL_KEY"): "/path/to/key",
            ("celery", "SSL_CERT"): "/path/to/cert",
            ("celery", "SSL_CACERT"): "",
        }
    )
    def test_ssl_active_without_cacert_uses_system_cas(self):
        """Test that empty SSL_CACERT falls back to system CAs (ca_certs omitted from config)."""
        import importlib
        import ssl

        importlib.reload(default_celery)
        broker_ssl = default_celery.DEFAULT_CELERY_CONFIG["broker_use_ssl"]

        assert "ca_certs" not in broker_ssl
        assert broker_ssl["cert_reqs"] == ssl.CERT_REQUIRED

    @conf_vars(
        {
            ("celery", "BROKER_URL"): "amqps://guest:guest@rabbitmq:5671//",
            ("celery", "SSL_ACTIVE"): "False",
        }
    )
    def test_amqps_broker_url_no_ssl_when_inactive(self):
        """Test that amqps:// broker URLs don't get SSL config when SSL_ACTIVE is False."""
        import importlib

        importlib.reload(default_celery)

        config = default_celery.DEFAULT_CELERY_CONFIG
        assert "broker_use_ssl" not in config

    @conf_vars(
        {
            ("celery", "BROKER_URL"): "amqps://guest:guest@rabbitmq:5671//",
            ("celery", "SSL_ACTIVE"): "True",
            ("celery", "SSL_MUTUAL_TLS"): "False",
            ("celery", "SSL_CACERT"): "/path/to/ca.pem",
        }
    )
    def test_amqps_one_way_tls(self):
        """Test one-way TLS for AMQP: only ca_certs, no keyfile/certfile."""
        import importlib
        import ssl

        importlib.reload(default_celery)

        config = default_celery.DEFAULT_CELERY_CONFIG
        assert "broker_use_ssl" in config
        broker_ssl = config["broker_use_ssl"]
        assert broker_ssl["ca_certs"] == "/path/to/ca.pem"
        assert broker_ssl["cert_reqs"] == ssl.CERT_REQUIRED
        assert "keyfile" not in broker_ssl
        assert "certfile" not in broker_ssl

    @conf_vars(
        {
            ("celery", "BROKER_URL"): "rediss://redis:6380//",
            ("celery", "SSL_ACTIVE"): "True",
            ("celery", "SSL_MUTUAL_TLS"): "False",
            ("celery", "SSL_CACERT"): "/path/to/ca.pem",
        }
    )
    def test_redis_one_way_tls(self):
        """Test one-way TLS for Redis: only ssl_ca_certs, no ssl_keyfile/ssl_certfile."""
        import importlib
        import ssl

        importlib.reload(default_celery)

        config = default_celery.DEFAULT_CELERY_CONFIG
        assert "broker_use_ssl" in config
        broker_ssl = config["broker_use_ssl"]
        assert broker_ssl["ssl_ca_certs"] == "/path/to/ca.pem"
        assert broker_ssl["ssl_cert_reqs"] == ssl.CERT_REQUIRED
        assert "ssl_keyfile" not in broker_ssl
        assert "ssl_certfile" not in broker_ssl

    @conf_vars(
        {
            ("celery", "BROKER_URL"): "amqps://guest:guest@rabbitmq:5671//",
            ("celery", "SSL_ACTIVE"): "True",
            ("celery", "SSL_MUTUAL_TLS"): "False",
            ("celery", "SSL_KEY"): "/path/to/key.pem",
            ("celery", "SSL_CERT"): "/path/to/cert.pem",
            ("celery", "SSL_CACERT"): "/path/to/ca.pem",
        }
    )
    def test_one_way_tls_ignores_key_cert(self):
        """Test that SSL_KEY/SSL_CERT are ignored when SSL_MUTUAL_TLS is False."""
        import importlib
        import ssl

        importlib.reload(default_celery)

        config = default_celery.DEFAULT_CELERY_CONFIG
        assert "broker_use_ssl" in config
        broker_ssl = config["broker_use_ssl"]
        assert broker_ssl["ca_certs"] == "/path/to/ca.pem"
        assert broker_ssl["cert_reqs"] == ssl.CERT_REQUIRED
        assert "keyfile" not in broker_ssl
        assert "certfile" not in broker_ssl


class TestCreateCeleryAppTeamIsolation:
    """Tests for create_celery_app() multi-team config isolation."""

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="ExecutorConf requires Airflow 3.2+")
    def test_custom_celery_config_options_applied(self):
        """User-provided celery_config_options (non-default) should be merged into team config."""
        custom_config = {"worker_concurrency": 42, "broker_url": "redis://custom:6379/0"}
        custom_path = "my_custom_module.CELERY_CONFIG"

        team_conf = ExecutorConf(team_name="team_alpha")
        original_get = team_conf.get

        def mock_get(section, key, **kwargs):
            if section == "celery" and key == "celery_config_options":
                return custom_path
            return original_get(section, key, **kwargs)

        mock_module = mock.MagicMock()
        mock_module.CELERY_CONFIG = custom_config

        with (
            mock.patch.object(team_conf, "get", side_effect=mock_get),
            mock.patch.object(celery_executor_utils, "import_module", return_value=mock_module),
        ):
            celery_app = celery_executor_utils.create_celery_app(team_conf)
            assert celery_app.conf.worker_concurrency == 42
            assert celery_app.conf.broker_url == "redis://custom:6379/0"

    def test_default_celery_config_options_skipped_via_identity_check(self):
        """When celery_config_options resolves to DEFAULT_CELERY_CONFIG (same object),
        it must be skipped — re-applying it would overwrite team-specific config
        since DEFAULT_CELERY_CONFIG is built from global conf."""
        original_get = conf.get
        # Path just needs a dot for rpartition and attr name matching DEFAULT_CELERY_CONFIG.
        # import_module is mocked to return default_celery module regardless of path.
        celery_config_path = "any.module.DEFAULT_CELERY_CONFIG"

        def mock_get(section, key, **kwargs):
            if section == "celery" and key == "celery_config_options":
                return celery_config_path
            return original_get(section, key, **kwargs)

        with (
            mock.patch.object(conf, "get", side_effect=mock_get),
            mock.patch.object(celery_executor_utils, "import_module") as mock_import,
        ):
            mock_import.return_value = default_celery
            celery_app = celery_executor_utils.create_celery_app(conf)
            # import_module called (path is non-None), but override skipped (same object)
            mock_import.assert_called_once()
            default_config = default_celery.get_default_celery_config(conf)
            assert celery_app.conf.broker_url == default_config["broker_url"]

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="ExecutorConf requires Airflow 3.2+")
    def test_team_specific_broker_not_overwritten(self):
        """Team-specific BROKER_URL set via ExecutorConf must survive create_celery_app()."""
        team_conf = ExecutorConf(team_name="team_alpha")

        original_get = team_conf.get

        def mock_team_get(section, key, **kwargs):
            if section == "celery" and key == "BROKER_URL":
                return "amqps://team-alpha-rabbit:5671//"
            return original_get(section, key, **kwargs)

        with mock.patch.object(team_conf, "get", side_effect=mock_team_get):
            celery_app = celery_executor_utils.create_celery_app(team_conf)
            assert celery_app.conf.broker_url == "amqps://team-alpha-rabbit:5671//"

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="ExecutorConf requires Airflow 3.2+")
    def test_team_app_name_includes_team_name(self):
        """Each team gets a unique Celery app name for broker isolation."""
        team_conf = ExecutorConf(team_name="team_beta")
        celery_app = celery_executor_utils.create_celery_app(team_conf)
        assert "team_beta" in celery_app.main
