airflow-pytest-operator 0.1.0__py3-none-any.whl

airflow_pytest_operator/__init__.py

@@ -0,0 +1,70 @@
+"""Run pytest suites as Airflow tasks.
+
+Public API:
+    PytestOperator         — the operator to use in DAGs
+    PytestRunner           — runner interface (extend for docker/k8s)
+    SubprocessPytestRunner — default runner
+    ResultParser           — parser interface
+    JUnitResultParser      — default parser
+    TestRunResult          — structured result model
+"""
+
+# Copyright 2026 Ilya Krysanov
+#
+# Licensed 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 typing import Any
+
+from .exceptions import (
+    AirflowPytestError,
+    ReportParseError,
+    TestExecutionError,
+    TestsFailedError,
+)
+from .models import CaseResult, RunArtifacts, TestRunResult
+from .operators import PytestOperator
+from .reporters import JUnitResultParser, ResultParser
+from .runners import PytestRunner, SubprocessPytestRunner
+
+__version__ = "0.1.0"
+
+
+def get_provider_info() -> dict[str, Any]:
+    """Metadata for Airflow's provider-discovery mechanism.
+
+    Lets Airflow's CLI/UI list this package as a provider. Optional —
+    operators work via plain imports regardless — but it makes the
+    package a well-behaved citizen if published.
+    """
+    return {
+        "package-name": "airflow-pytest-operator",
+        "name": "Pytest Operator",
+        "description": "Run pytest suites as Airflow tasks.",
+        "versions": [__version__],
+    }
+
+
+__all__ = [
+    "PytestOperator",
+    "PytestRunner",
+    "SubprocessPytestRunner",
+    "ResultParser",
+    "JUnitResultParser",
+    "TestRunResult",
+    "RunArtifacts",
+    "CaseResult",
+    "AirflowPytestError",
+    "TestExecutionError",
+    "ReportParseError",
+    "TestsFailedError",
+]

airflow_pytest_operator/compat/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2026 Ilya Krysanov
+#
+# Licensed 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 .airflow import BaseOperator, apply_defaults, get_airflow_version
+
+__all__ = ["BaseOperator", "apply_defaults", "get_airflow_version"]

airflow_pytest_operator/compat/airflow.py

@@ -0,0 +1,101 @@
+"""Airflow version compatibility shim.
+
+This is the *only* module in the package that imports Airflow directly.
+Everything else imports ``BaseOperator`` from here. Centralizing the
+version-specific imports means that supporting a new Airflow release is
+a one-file change (Open/Closed at the package level).
+
+Airflow 2.x and 3.x differ in the import path of ``BaseOperator`` and a
+few helpers. We resolve them once, lazily, and expose a stable surface.
+"""
+
+# Copyright 2026 Ilya Krysanov
+#
+# Licensed 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
+
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    # Static analysers (mypy/Pylance) don't have Airflow installed and
+    # can't follow the runtime try/except below. We give them a concrete
+    # but minimal type to bind ``BaseOperator`` to, so downstream code is
+    # type-checkable. At runtime this block is skipped entirely.
+    class BaseOperator:  # noqa: D101 - stub for type-checking only
+        task_id: str
+        log: Any
+
+        def __init__(self, *args: Any, **kwargs: Any) -> None: ...
+
+
+@lru_cache(maxsize=1)
+def get_airflow_version() -> tuple[int, ...]:
+    try:
+        from airflow.version import version as _v
+    except Exception:  # pragma: no cover - airflow always ships this
+        return (0,)
+    parts: list[int] = []
+    for chunk in _v.split(".")[:3]:
+        num = "".join(ch for ch in chunk if ch.isdigit())
+        parts.append(int(num) if num else 0)
+    return tuple(parts)
+
+
+def _import_base_operator() -> type[Any]:
+    """Return the correct BaseOperator class for the installed Airflow.
+
+    Airflow 3 moved BaseOperator to ``airflow.sdk``. Airflow 2 keeps it at
+    ``airflow.models.baseoperator``. We try the new location first and
+    fall back, so a single wheel works against both.
+    """
+    try:
+        # Airflow 3.x (Task SDK)
+        from airflow.sdk import BaseOperator  # type: ignore[attr-defined]
+
+        return BaseOperator  # type: ignore[no-any-return]
+    except Exception:
+        pass
+    from airflow.models.baseoperator import BaseOperator  # Airflow 2.x
+
+    return BaseOperator  # type: ignore[no-any-return]
+
+
+def _import_apply_defaults() -> Any:
+    """``apply_defaults`` is a no-op decorator in 2.x and gone in 3.x.
+
+    We return a passthrough when it's unavailable so operator code can
+    reference it uniformly without branching.
+    """
+    try:
+        from airflow.utils.decorators import (  # type: ignore[attr-defined]
+            apply_defaults,
+        )
+
+        return apply_defaults
+    except Exception:
+
+        def apply_defaults(func: Any) -> Any:
+            return func
+
+        return apply_defaults
+
+
+# Resolved at import time, but cheap and side-effect-free.
+# The TYPE_CHECKING stub above already bound ``BaseOperator`` for analysers;
+# this is the real runtime value, hence the explicit no-redef suppression.
+BaseOperator = _import_base_operator()  # type: ignore[assignment,no-redef,misc]
+apply_defaults = _import_apply_defaults()
+
+__all__ = ["BaseOperator", "apply_defaults", "get_airflow_version"]

airflow_pytest_operator/exceptions.py

@@ -0,0 +1,58 @@
+"""Exception hierarchy for the operator.
+
+A small, focused hierarchy lets callers (and Airflow's retry logic)
+distinguish *test failures* from *infrastructure failures*. That
+distinction matters: a failing test usually shouldn't be retried,
+but a missing pytest binary or unreadable report might be.
+"""
+
+# Copyright 2026 Ilya Krysanov
+#
+# Licensed 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
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .models import TestRunResult
+
+
+class AirflowPytestError(Exception):
+    """Base class for all errors raised by this package."""
+
+
+class TestExecutionError(AirflowPytestError):
+    """The runner could not execute pytest at all (binary missing, etc.)."""
+
+    __test__ = False
+
+
+class ReportParseError(AirflowPytestError):
+    """A report file was produced but could not be parsed."""
+
+
+class TestsFailedError(AirflowPytestError):
+    """Pytest ran successfully but one or more tests failed.
+
+    Carries the structured result so downstream handlers can inspect it.
+    """
+
+    __test__ = False
+
+    def __init__(self, result: TestRunResult) -> None:
+        self.result = result
+        super().__init__(
+            f"{result.failed} failed, {result.errors} errors "
+            f"out of {result.total} tests"
+        )

airflow_pytest_operator/models.py

@@ -0,0 +1,95 @@
+"""Domain models for test runs.
+
+These are plain, framework-agnostic dataclasses. Nothing here imports
+Airflow, pytest, or subprocess — keeping the domain layer dependency-free
+makes it trivial to unit-test parsers and the operator in isolation (DIP).
+"""
+
+# Copyright 2026 Ilya Krysanov
+#
+# Licensed 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
+
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+
+@dataclass(frozen=True)
+class CaseResult:
+    """A single test case outcome parsed from a report."""
+
+    name: str
+    classname: str
+    time: float
+    outcome: str  # "passed" | "failed" | "error" | "skipped"
+    message: str | None = None
+
+    @property
+    def node_id(self) -> str:
+        """Reconstruct a pytest-style node id when possible."""
+        if self.classname:
+            return f"{self.classname}::{self.name}"
+        return self.name
+
+
+@dataclass(frozen=True)
+class RunArtifacts:
+    """Everything a Runner produces: where to find outputs + the exit code.
+
+    A Runner's job ends at producing files; interpreting them is the
+    Parser's job. This separation is what lets us swap runners
+    (subprocess, docker, k8s-pod) without touching parsing logic.
+    """
+
+    exit_code: int
+    junit_xml_path: str | None
+    stdout: str = ""
+    stderr: str = ""
+    working_dir: str | None = None
+
+
+@dataclass(frozen=True)
+class TestRunResult:
+    """Structured, serializable result of a pytest run."""
+
+    __test__ = False  # not a pytest test class despite the name
+
+    total: int
+    passed: int
+    failed: int
+    skipped: int
+    errors: int
+    duration: float
+    exit_code: int
+    cases: list[CaseResult] = field(default_factory=list)
+
+    @property
+    def success(self) -> bool:
+        return self.failed == 0 and self.errors == 0
+
+    @property
+    def failed_node_ids(self) -> list[str]:
+        return [c.node_id for c in self.cases if c.outcome in ("failed", "error")]
+
+    def to_xcom(self) -> dict[str, Any]:
+        """A compact, JSON-serializable dict suitable for XCom.
+
+        We deliberately drop per-case ``message`` blobs from the summary
+        pushed to XCom (they can be huge); full detail stays in logs.
+        """
+        data = asdict(self)
+        data.pop("cases", None)
+        data["success"] = self.success
+        data["failed_node_ids"] = self.failed_node_ids
+        return data

airflow_pytest_operator/operators/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2026 Ilya Krysanov
+#
+# Licensed 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 .pytest_operator import PytestOperator
+
+__all__ = ["PytestOperator"]

airflow_pytest_operator/operators/pytest_operator.py

@@ -0,0 +1,192 @@
+"""The Airflow operator.
+
+This class is intentionally *thin*. Its only responsibilities are:
+  1. orchestrate runner -> parser,
+  2. integrate with Airflow (templating, XCom, logging, fail policy).
+
+It contains no subprocess logic and no XML parsing. Both collaborators
+are injected (Dependency Inversion): defaults are provided for ergonomic
+use in DAGs, but tests can pass fakes, and advanced users can swap in a
+Docker/K8s runner or a JSON parser without subclassing.
+"""
+
+# Copyright 2026 Ilya Krysanov
+#
+# Licensed 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
+
+from collections.abc import Sequence
+from typing import Any
+
+from ..compat import BaseOperator
+from ..exceptions import TestExecutionError, TestsFailedError
+from ..reporters import JUnitResultParser, ResultParser
+from ..runners import PytestRunner, SubprocessPytestRunner
+
+
+class PytestOperator(BaseOperator):
+    """Run a pytest suite as an Airflow task.
+
+    :param test_path: file or directory to pass to pytest (templated).
+    :param pytest_args: extra CLI args, e.g. ``["-k", "smoke", "-x"]`` (templated).
+    :param env: extra environment variables for the run (templated).
+    :param fail_on_test_failure: if True (default) the task fails when any
+        test fails or errors; if False the task always succeeds and the
+        outcome is only reflected in XCom.
+    :param push_result: if True (default) push the structured summary dict
+        to XCom under the ``pytest_result`` key. If False, nothing is sent
+        to XCom at all -- this also disables Airflow's automatic push of
+        the task's return value (``do_xcom_push`` is forced off).
+    :param runner: injectable :class:`PytestRunner` (default: subprocess).
+    :param parser: injectable :class:`ResultParser` (default: JUnit).
+    """
+
+    # Airflow Jinja-templates these attributes before execute() runs.
+    template_fields: Sequence[str] = ("test_path", "pytest_args", "env")
+    ui_color = "#4caf50"
+
+    def __init__(
+        self,
+        *,
+        test_path: str,
+        pytest_args: Sequence[str] | None = None,
+        env: dict[str, str] | None = None,
+        fail_on_test_failure: bool = True,
+        push_result: bool = True,
+        runner: PytestRunner | None = None,
+        parser: ResultParser | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.test_path = test_path
+        self.pytest_args = list(pytest_args) if pytest_args else []
+        self.env = env or {}
+        self.fail_on_test_failure = fail_on_test_failure
+        self.push_result = push_result
+        if not push_result:
+            # push_result=False means "nothing in XCom at all". Airflow
+            # otherwise auto-pushes execute()'s return value under the
+            # "return_value" key when do_xcom_push is True (the default),
+            # so we disable that here to honour the intent. This wins over
+            # an explicit do_xcom_push=True in kwargs by design: choosing
+            # push_result=False is the more specific, intentional signal.
+            self.do_xcom_push = False
+        # DI with sensible defaults — collaborators, not inheritance.
+        self._runner = runner or SubprocessPytestRunner()
+        self._parser = parser or JUnitResultParser()
+
+    def execute(self, context: Any) -> dict[str, Any] | None:
+        self.log.info("Running pytest on %s", self.test_path)
+
+        artifacts = self._runner.run(
+            self.test_path,
+            pytest_args=self.pytest_args,
+            env=self.env,
+        )
+
+        # Surface child output in the task log regardless of outcome.
+        if artifacts.stdout:
+            self.log.info("pytest stdout:\n%s", artifacts.stdout)
+        if artifacts.stderr:
+            self.log.warning("pytest stderr:\n%s", artifacts.stderr)
+
+        # ``run_ok`` drives the runner's cleanup policy ("on_success").
+        # It stays False until we've parsed a report whose tests all passed,
+        # so any early exit (missing report, failing tests) is treated as a
+        # failure and artifacts can be retained for post-mortem.
+        run_ok = False
+        try:
+            # No report means pytest never got far enough to write one
+            # (collection error, internal crash, OOM kill, wrong path).
+            # This is an *execution* failure, not a test failure -- surface
+            # it clearly with the captured stderr, not a cryptic parse error.
+            if artifacts.junit_xml_path is None:
+                raise TestExecutionError(
+                    "pytest produced no JUnit report "
+                    f"(exit code {artifacts.exit_code}). "
+                    "This usually means a collection error or crash before "
+                    "any test ran. Captured stderr:\n"
+                    f"{artifacts.stderr or '<empty>'}"
+                )
+
+            result = self._parser.parse(
+                artifacts.junit_xml_path, exit_code=artifacts.exit_code
+            )
+
+            self.log.info(
+                "Results: total=%d passed=%d failed=%d errors=%d skipped=%d (%.2fs)",
+                result.total,
+                result.passed,
+                result.failed,
+                result.errors,
+                result.skipped,
+                result.duration,
+            )
+            if result.failed_node_ids:
+                self.log.error(
+                    "Failed tests:\n  %s", "\n  ".join(result.failed_node_ids)
+                )
+
+            summary = result.to_xcom() if self.push_result else None
+            if self.push_result:
+                # Returning a value auto-pushes to XCom under the default
+                # key; we also push an explicit, named key for stable
+                # downstream references.
+                context["ti"].xcom_push(key="pytest_result", value=summary)
+
+            run_ok = result.success
+            if self.fail_on_test_failure and not result.success:
+                raise TestsFailedError(result)
+
+            return summary
+        finally:
+            # Always invoke cleanup; the runner decides what to remove
+            # based on its policy and the success flag. Never let cleanup
+            # errors mask the real outcome of execute().
+            try:
+                self._runner.cleanup(success=run_ok)
+            except Exception:  # pragma: no cover - best-effort teardown
+                self.log.exception("Error while cleaning up report directory")
+
+    def on_kill(self) -> None:
+        """Abort the test run when Airflow terminates the task.
+
+        Airflow calls this when the task is killed -- execution timeout,
+        a manual clear/mark-failed, or the worker shutting down (SIGTERM).
+        We delegate to the runner, which owns the actual process/resource;
+        the operator deliberately knows nothing about subprocesses.
+
+        Delegation keeps responsibilities separate: the operator handles
+        the Airflow lifecycle, the runner handles teardown of whatever it
+        spawned. Runners that have nothing to cancel inherit a safe no-op.
+        """
+        self.log.warning("Task killed -- cancelling pytest run on %s", self.test_path)
+        try:
+            self._runner.cancel()
+        except Exception:  # pragma: no cover - best-effort teardown
+            # on_kill must never raise: it runs during teardown, and an
+            # exception here can mask the original termination cause and
+            # leave the task in a confusing state.
+            self.log.exception("Error while cancelling pytest run")
+
+        # A killed run is never successful, but with the default "always"
+        # policy the temp report dir is still removed -- kills/timeouts are
+        # exactly when leaked dirs pile up, so we must clean here too. The
+        # cancel() above has already stopped the process, so nothing is
+        # still writing into the directory. cleanup() is idempotent and
+        # thread-safe, so racing with execute()'s own finally is harmless.
+        try:
+            self._runner.cleanup(success=False)
+        except Exception:  # pragma: no cover - best-effort teardown
+            self.log.exception("Error while cleaning up report directory")

airflow_pytest_operator/reporters/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2026 Ilya Krysanov
+#
+# Licensed 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 .base import ResultParser
+from .junit_parser import JUnitResultParser
+
+__all__ = ["ResultParser", "JUnitResultParser"]

airflow_pytest_operator/reporters/base.py

@@ -0,0 +1,43 @@
+"""The result-parser interface.
+
+A parser turns a report file into a :class:`TestRunResult`. It knows
+nothing about how the report was produced. Keeping this separate from
+the runner means we can support other report formats (e.g. a JSON
+report plugin) by adding a parser, not by editing existing code (OCP).
+"""
+
+# Copyright 2026 Ilya Krysanov
+#
+# Licensed 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
+
+from abc import ABC, abstractmethod
+
+from ..models import TestRunResult
+
+
+class ResultParser(ABC):
+    """Parses a report file into a structured result."""
+
+    @abstractmethod
+    def parse(self, report_path: str, *, exit_code: int = 0) -> TestRunResult:
+        """Parse ``report_path`` into a :class:`TestRunResult`.
+
+        ``exit_code`` is threaded through so the result records how the
+        process actually terminated (a parser can't always infer e.g.
+        an internal pytest error from the XML alone).
+
+        Raises :class:`ReportParseError` if the file is missing or malformed.
+        """
+        raise NotImplementedError