System Tests¶

"""Example DAGs for apache-airflow-providers-pixi, run by test_example_dags.py through dag.test()."""

from airflow.sdk import DAG, task

from airflow.providers.pixi.operators.pixi import PixiOperator

ENV = {"dependencies": {"python": "3.12.*"}}

with DAG("pixi_example", params={"project": "", "out": ""}) as dag_ok:

    @task
    def produce() -> int:
        return 20

    # Defined in the DAG file, so it reaches the environment as source.
    @task.pixi(**ENV)
    def double(x: int) -> dict:
        import sys

        print("doubling", x)
        return {"value": x * 2, "prefix": sys.prefix}


    # A module of the project, whose path is templated; the first argument is double's XCom.
    PixiOperator(
        task_id="report",
        pixi_project_path="{{ params.project }}",
        python_callable="report:write",
        op_args=[double(produce()), "{{ params.out }}"],
    )


with DAG("pixi_example_fail") as dag_fail:

    @task.pixi(**ENV)
    def boom() -> None:
        raise ValueError("task failed on purpose")

    boom()

"""Example DAGs for the Bash operator, the sensors, the Airflow context and failing callables.

Run by test_example_bash_sensor.py through dag.test(). ``params.project`` is a pixi project, ``params.out`` a
directory the tasks write their results to.
"""

from __future__ import annotations

import os
import tempfile

from airflow.sdk import DAG, task

from airflow.providers.pixi.operators.bash import PixiBashOperator
from airflow.providers.pixi.operators.pixi import PixiOperator
from airflow.providers.pixi.sensors.pixi import PixiSensor

PARAMS = {"project": "", "out": ""}
ENV_CACHE = os.environ.get("PIXI_E2E_ENV_CACHE") or os.path.join(tempfile.gettempdir(), "airflow-pixi-e2e-envs")


def count_pokes(path: str, pokes: int) -> dict:
    """Record the environment of this poke in ``path``; done on poke number ``pokes``."""
    import sys

    with open(path, "a") as f:
        f.write(sys.prefix + "\n")
    with open(path) as f:
        done = len(f.read().splitlines())
    return {"is_done": done >= pokes, "xcom_value": done}


def context_values(label: str, ds=None, run_id=None, params=None, logical_date=None) -> dict:
    import json
    import os

    values = {"label": label, "ds": ds, "run_id": run_id, "params": params, "logical_date": logical_date}
    values["greeting"] = os.environ.get("GREETING")
    with open(os.path.join(params["out"], "context.json"), "w") as f:
        json.dump(values, f)
    return values


def fail_on_purpose(path: str) -> None:
    raise RuntimeError(f"{path} is empty")


def exit_with_99() -> None:
    import sys

    sys.exit(99)


with DAG("pixi_bash_sensor_example", params=PARAMS) as dag_bash_sensor:
    bash = PixiBashOperator(
        task_id="bash",
        pixi_project_path="{{ params.project }}",
        env={"GREETING": "hello"},
        append_env=True,
        bash_command="python -c 'import os, sys; print(os.environ[\"GREETING\"], sys.prefix)'",
    )

    @task.pixi_bash(pixi_project_path="{{ params.project }}", lock_mode="locked")
    def bash_from_function(word: str) -> str:
        return f"echo \"{word} {{{{ ds }}}} $(python -c 'import sys; print(sys.prefix)')\""

    poke = PixiSensor(
        task_id="poke",
        pixi_project_path="{{ params.project }}",
        lock_mode="frozen",
        python_callable=count_pokes,
        op_args=["{{ params.out }}/poke.txt", 2],
        mode="poke",
        poke_interval=1,
        timeout=600,
    )

    @task.pixi_sensor(
        dependencies={"python": "3.12.*"}, env_cache_path=ENV_CACHE, mode="reschedule", poke_interval=1, timeout=900
    )
    def reschedule(path: str, pokes: int) -> dict:
        import sys

        with open(path, "a") as f:
            f.write(sys.prefix + "\n")
        with open(path) as f:
            done = len(f.read().splitlines())
        return {"is_done": done >= pokes, "xcom_value": done}

    @task
    def collect(bash_line: str, function_line: str, pokes: int, reschedules: int, out: str) -> None:
        import json

        values = {"bash": bash_line, "bash_from_function": function_line, "poke": pokes, "reschedule": reschedules}
        with open(os.path.join(out, "collected.json"), "w") as f:
            json.dump(values, f)

    from_function = bash_from_function("rendered")
    rescheduled = reschedule("{{ params.out }}/reschedule.txt", 2)
    bash >> [from_function, poke]
    collect(bash.output, from_function, poke.output, rescheduled, "{{ params.out }}")


with DAG("pixi_context_example", params=PARAMS) as dag_context:
    PixiOperator(
        task_id="context_values",
        pixi_project_path="{{ params.project }}",
        python_callable=context_values,
        op_args=["from op_args"],
        env_from_variables={"GREETING": "pixi_e2e_greeting"},
    )


with DAG("pixi_callable_error_example", params=PARAMS) as dag_callable_error:
    PixiOperator(
        task_id="raises",
        pixi_project_path="{{ params.project }}",
        python_callable=fail_on_purpose,
        op_args=["input.csv"],
    )

    PixiOperator(
        task_id="skips",
        pixi_project_path="{{ params.project }}",
        python_callable=exit_with_99,
        skip_on_exit_code=99,
    )

"""Example DAGs for PixiKubernetesPodOperator, run in a kind cluster by test_example_kubernetes.py."""

import datetime
from pathlib import Path

from airflow.sdk import DAG, task
from kubernetes.client import models as k8s

from airflow.providers.pixi.operators.kubernetes import PixiKubernetesPodOperator

CACHE = k8s.V1Volume(
    name="cache", host_path=k8s.V1HostPathVolumeSource(path="/var/tmp/pixi-airflow", type="DirectoryOrCreate")
)
POD = {
    "namespace": "default",
    "in_cluster": False,
    "startup_timeout_seconds": 600,
    "dependencies": {"python": "3.12.*"},
    "env_cache_path": "/cache/envs",
    "env_vars": {"PIXI_CACHE_DIR": "/cache/pixi"},
    "volumes": [CACHE],
    "volume_mounts": [k8s.V1VolumeMount(name="cache", mount_path="/cache")],
}


def stamp():
    import datetime

    return datetime.datetime(2026, 1, 2, 3, 4, 5, tzinfo=datetime.timezone.utc)


with DAG("pixi_kubernetes_example") as dag_ok:

    @task.pixi_kubernetes(**POD)
    def summarize(values: list, run_id=None) -> dict:
        import statistics
        import sys

        return {"mean": statistics.mean(values), "run_id": run_id, "python": list(sys.version_info[:2])}

    pickled = PixiKubernetesPodOperator(task_id="pickled", python_callable=stamp, serializer="pickle", **POD)
    deferred = PixiKubernetesPodOperator(
        task_id="deferred", python_callable=stamp, serializer="pickle", deferrable=True, poll_interval=5, **POD
    )
    no_xcom = PixiKubernetesPodOperator(
        task_id="no_xcom", python_callable="builtins:max", op_args=[1, 2], do_xcom_push=False, **POD
    )

    @task
    def check(summary: dict, pickled: datetime.datetime, deferred: datetime.datetime, run_id=None, ti=None) -> None:
        expected = stamp()
        assert summary == {"mean": 2, "run_id": run_id, "python": [3, 12]}, summary
        assert pickled == expected, pickled
        assert deferred == expected, deferred
        assert ti.xcom_pull(task_ids="no_xcom") is None

    no_xcom >> check(summarize([1, 2, 3]), pickled.output, deferred.output)


def record_exception(context) -> None:
    exception = context["exception"]
    out = Path(context["params"]["out"]) / f"{context['ti'].task_id}.txt"
    out.write_text(f"{type(exception).__name__}: {exception}")


def boom(rows: str) -> None:
    raise ValueError(f"no rows in {rows}")


with DAG("pixi_kubernetes_example_fail", params={"out": ""}) as dag_fail:
    failing_pod = {**POD, "python_callable": boom, "op_args": ["sales"], "on_failure_callback": record_exception}
    PixiKubernetesPodOperator(task_id="boom", **failing_pod)
    PixiKubernetesPodOperator(task_id="boom_deferred", deferrable=True, poll_interval=5, **failing_pod)