Repository: apache/airflow
File: helm-tests/tests/helm_tests/airflow_core/test_worker.py (Python)

Code before:
            ("2.1.0", "airflow celery worker"),
        ],
    )
    def test_default_command_and_args_airflow_version(self, airflow_version, expected_arg):
        docs = render_chart(
            values={
                "airflowVersion": airflow_version,
            },
            show_only=["templates/workers/worker-deployment.yaml"],
        )

        assert jmespath.search("spec.template.spec.containers[0].command", docs[0]) is None
        assert [
            "bash",
            "-c",
            f"exec \\\n{expected_arg}",
        ] == jmespath.search("spec.template.spec.containers[0].args", docs[0])

    @pytest.mark.parametrize(
        ("workers_values", "expected"),
        [
            ({"command": ["custom", "command"]}, ["custom", "command"]),
            ({"command": ["custom", "{{ .Release.Name }}"]}, ["custom", "release-name"]),
            ({"celery": {"command": ["custom", "command"]}}, ["custom", "command"]),
            ({"celery": {"command": ["custom", "{{ .Release.Name }}"]}}, ["custom", "release-name"]),
            ({"command": ["test"], "celery": {"command": ["custom", "command"]}}, ["custom", "command"]),
        ],
    )
    def test_should_add_command(self, workers_values, expected):
        docs = render_chart(
            values={"workers": workers_values},
            show_only=["templates/workers/worker-deployment.yaml"],
        )

        assert expected == jmespath.search("spec.template.spec.containers[0].command", docs[0])

    @pytest.mark.parametrize(
        "workers_values",
        [
            {"command": None},
            {"command": []},
            {"celery": {"command": None}},
            {"celery": {"command": []}},
        ],
    )
    def test_should_not_add_command(self, workers_values):
        docs = render_chart(
            values={"workers": workers_values},
            show_only=["templates/workers/worker-deployment.yaml"],
        )


Review comment (suggestion):
Personally, I would split tests to test particular arguments and this one left only for testing `enabled` flag

Code after:
                "kerberosInitContainer": {"enabled": False},
                "celery": {"kerberosInitContainer": {"enabled": True}},
            },
        ],
    )
    def test_airflow_kerberos_init_container_celery_values(self, workers_values):
        """Test that workers.celery.kerberosInitContainer configuration works and takes precedence."""
        docs = render_chart(
            values={
                "airflowVersion": "2.8.0",
                "workers": {
                    **workers_values,
                    "celery": {
                        **workers_values.get("celery", {}),
                        "persistence": {"fixPermissions": True},
                    },
                },
            },
            show_only=["templates/workers/worker-deployment.yaml"],
        )

        initContainers = jmespath.search("spec.template.spec.initContainers", docs[0])
        # Should have 3 init containers: wait-for-migrations, kerberos-init, volume-permissions
        assert len(initContainers) == 3
        assert initContainers[1]["name"] == "kerberos-init"
        assert initContainers[1]["args"] == ["kerberos", "-o"]

    def test_airflow_kerberos_init_container_resources(self):
        """Test that kerberos init container resources can be configured via workers.celery.kerberosInitContainer."""
        docs = render_chart(
            values={
                "airflowVersion": "2.8.0",
                "workers": {
                    "celery": {
                        "kerberosInitContainer": {
                            "enabled": True,
                            "resources": {
                                "limits": {"cpu": "100m", "memory": "128Mi"},
                                "requests": {"cpu": "50m", "memory": "64Mi"},
                            },
                        },
                    },
                },
            },
            show_only=["templates/workers/worker-deployment.yaml"],
        )

        initContainers = jmespath.search("spec.template.spec.initContainers", docs[0])
        kerberos_init = next((c for c in initContainers if c["name"] == "kerberos-init"), None)
        assert kerberos_init is not None
        assert kerberos_init["resources"]["limits"]["cpu"] == "100m"