ariadne/tests/test_k8s_jobs.py

from __future__ import annotations

import pytest

from ariadne.k8s.jobs import JobSpawner


def test_job_from_cronjob_applies_env_and_ttl() -> None:
    cronjob = {
        "spec": {
            "jobTemplate": {
                "spec": {
                    "template": {
                        "spec": {
                            "containers": [
                                {"name": "sync", "env": [{"name": "FOO", "value": "1"}]}
                            ]
                        }
                    }
                }
            }
        }
    }

    spawner = JobSpawner("ns", "cron")
    job = spawner._job_from_cronjob(
        cronjob,
        "User@Name",
        env_overrides=[{"name": "FOO", "value": "2"}, {"name": "BAR", "value": "3"}],
        job_ttl_seconds=3600,
    )

    assert job["spec"]["ttlSecondsAfterFinished"] == 3600
    labels = job["metadata"]["labels"]
    assert labels["atlas.bstein.dev/trigger"] == "ariadne"

    env = job["spec"]["template"]["spec"]["containers"][0]["env"]
    env_map = {item["name"]: item["value"] for item in env}
    assert env_map["FOO"] == "2"
    assert env_map["BAR"] == "3"


def test_job_from_cronjob_env_not_list() -> None:
    cronjob = {
        "spec": {
            "jobTemplate": {
                "spec": {
                    "template": {
                        "spec": {
                            "containers": [
                                {"name": "sync", "env": "bad"}
                            ]
                        }
                    }
                }
            }
        }
    }

    spawner = JobSpawner("ns", "cron")
    job = spawner._job_from_cronjob(
        cronjob,
        "label",
        env_overrides=[{"name": "FOO", "value": "1"}],
    )
    env = job["spec"]["template"]["spec"]["containers"][0]["env"]
    assert env == [{"name": "FOO", "value": "1"}]


def test_safe_name_fragment() -> None:
    assert JobSpawner._safe_name_fragment("User@Name") == "user-name"
    assert JobSpawner._safe_name_fragment("$$$") == "job"


def test_trigger_creates_job(monkeypatch) -> None:
    cronjob = {
        "metadata": {"name": "cron"},
        "spec": {
            "jobTemplate": {
                "spec": {
                    "template": {
                        "spec": {
                            "containers": [{"name": "sync", "env": []}]
                        }
                    }
                }
            }
        },
    }

    monkeypatch.setattr("ariadne.k8s.jobs.get_json", lambda *args, **kwargs: cronjob)
    monkeypatch.setattr("ariadne.k8s.jobs.post_json", lambda *args, **kwargs: {"metadata": {"name": "job"}})

    spawner = JobSpawner("ns", "cron")
    result = spawner.trigger("label", None, job_ttl_seconds=30)

    assert result["status"] == "queued"


def test_trigger_uses_manifest(monkeypatch) -> None:
    cronjob = {
        "metadata": {"name": "cron"},
        "spec": {
            "jobTemplate": {
                "spec": {
                    "template": {
                        "spec": {
                            "containers": [{"name": "sync", "env": []}]
                        }
                    }
                }
            }
        },
    }

    def explode(*_args, **_kwargs):
        raise AssertionError("get_json should not be called")

    monkeypatch.setattr("ariadne.k8s.jobs.get_json", explode)
    monkeypatch.setattr("ariadne.k8s.jobs.post_json", lambda *args, **kwargs: {"metadata": {"name": "job"}})

    spawner = JobSpawner("ns", "cron", manifest=cronjob)
    result = spawner.trigger("label", None, job_ttl_seconds=30)

    assert result["status"] == "queued"


def test_trigger_missing_job_name(monkeypatch) -> None:
    cronjob = {"spec": {"jobTemplate": {"spec": {"template": {"spec": {"containers": [{"name": "sync"}]}}}}}}
    monkeypatch.setattr("ariadne.k8s.jobs.get_json", lambda *args, **kwargs: cronjob)
    posted = {}

    def fake_post(_path, payload):
        posted["payload"] = payload
        return {}

    monkeypatch.setattr("ariadne.k8s.jobs.post_json", fake_post)

    spawner = JobSpawner("ns", "cron")
    result = spawner.trigger("label", None)
    assert result["job"] == posted["payload"]["metadata"]["name"]


def test_trigger_missing_job_name_raises(monkeypatch) -> None:
    monkeypatch.setattr("ariadne.k8s.jobs.get_json", lambda *args, **kwargs: {})
    monkeypatch.setattr("ariadne.k8s.jobs.post_json", lambda *args, **kwargs: {})
    spawner = JobSpawner("ns", "cron")
    monkeypatch.setattr(spawner, "_job_from_cronjob", lambda *args, **kwargs: {"metadata": {}})
    with pytest.raises(RuntimeError):
        spawner.trigger("label", None)


def test_wait_for_completion_success(monkeypatch) -> None:
    responses = [
        {"status": {"succeeded": 1}},
    ]

    def fake_get_json(path):
        return responses.pop(0)

    monkeypatch.setattr("ariadne.k8s.jobs.get_json", fake_get_json)

    spawner = JobSpawner("ns", "cron")
    result = spawner.wait_for_completion("job", timeout_sec=0.1)
    assert result["status"] == "ok"


def test_wait_for_completion_skips_bad_condition(monkeypatch) -> None:
    responses = [
        {"status": {"conditions": ["bad", {"type": "Complete", "status": "True"}]}},
    ]
    monkeypatch.setattr("ariadne.k8s.jobs.get_json", lambda path: responses.pop(0))

    spawner = JobSpawner("ns", "cron")
    result = spawner.wait_for_completion("job", timeout_sec=0.1)
    assert result["status"] == "ok"


def test_wait_for_completion_error(monkeypatch) -> None:
    responses = [
        {"status": {"failed": 1}},
    ]

    monkeypatch.setattr("ariadne.k8s.jobs.get_json", lambda path: responses.pop(0))

    spawner = JobSpawner("ns", "cron")
    result = spawner.wait_for_completion("job", timeout_sec=0.1)
    assert result["status"] == "error"


def test_wait_for_completion_timeout(monkeypatch) -> None:
    monkeypatch.setattr("ariadne.k8s.jobs.get_json", lambda path: {"status": {}})

    spawner = JobSpawner("ns", "cron")
    result = spawner.wait_for_completion("job", timeout_sec=0.01)
    assert result["status"] == "running"


def test_wait_for_completion_conditions(monkeypatch) -> None:
    responses = [
        {"status": {"conditions": [{"type": "Complete", "status": "True"}]}},
    ]
    monkeypatch.setattr("ariadne.k8s.jobs.get_json", lambda path: responses.pop(0))

    spawner = JobSpawner("ns", "cron")
    result = spawner.wait_for_completion("job", timeout_sec=0.1)
    assert result["status"] == "ok"


def test_wait_for_completion_failed_condition(monkeypatch) -> None:
    responses = [
        {"status": {"conditions": [{"type": "Failed", "status": "True"}]}},
    ]
    monkeypatch.setattr("ariadne.k8s.jobs.get_json", lambda path: responses.pop(0))

    spawner = JobSpawner("ns", "cron")
    result = spawner.wait_for_completion("job", timeout_sec=0.1)
    assert result["status"] == "error"


def test_trigger_and_wait(monkeypatch) -> None:
    cronjob = {"spec": {"jobTemplate": {"spec": {"template": {"spec": {"containers": [{"name": "sync"}]}}}}}}
    monkeypatch.setattr("ariadne.k8s.jobs.get_json", lambda *args, **kwargs: cronjob)
    monkeypatch.setattr(
        "ariadne.k8s.jobs.post_json",
        lambda *args, **kwargs: {"metadata": {"name": "job"}},
    )
    monkeypatch.setattr(
        "ariadne.k8s.jobs.JobSpawner.wait_for_completion",
        lambda self, job, timeout_sec: {"job": job, "status": "ok"},
    )

    spawner = JobSpawner("ns", "cron")
    result = spawner.trigger_and_wait("label", None, timeout_sec=1.0)
    assert result["status"] == "ok"


def test_trigger_and_wait_missing_job_name(monkeypatch) -> None:
    spawner = JobSpawner("ns", "cron")
    monkeypatch.setattr(spawner, "trigger", lambda *args, **kwargs: {"job": ""})
    with pytest.raises(RuntimeError):
        spawner.trigger_and_wait("label", None, timeout_sec=0.1)