How to Mock a Kubeflow Pipeline in Python

from typing import Dict, Any
from dataclasses import dataclass, field
from collections import OrderedDict


@dataclass
class KubeflowPipelineMock:
    """A minimal mock of a Kubeflow pipeline DAG."""
    name: str
    components: OrderedDict[str, callable] = field(default_factory=OrderedDict)

    def add_component(self, name: str, func: callable) -> "KubeflowPipelineMock":
        """Register a component (function) in the pipeline."""
        self.components[name] = func
        return self

    def execute(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Run the pipeline, passing outputs as inputs to the next stage."""
        data = dict(inputs)
        for comp_name, comp_func in self.components.items():
            result = comp_func(data)
            data.update(result)
            print(f"[{comp_name}] -> {list(result.keys())}")
        return data


def load_step(context: Dict[str, Any]) -> Dict[str, Any]:
    return {"loaded": context.get("raw", "default")}


def transform_step(context: Dict[str, Any]) -> Dict[str, Any]:
    return {"transformed": context["loaded"].upper()}


def save_step(context: Dict[str, Any]) -> Dict[str, Any]:
    return {"saved": f"final: {context['transformed']}"}


if __name__ == "__main__":
    pipeline = (
        KubeflowPipelineMock(name="sample-dag")
        .add_component("load", load_step)
        .add_component("transform", transform_step)
        .add_component("save", save_step)
    )
    outputs = pipeline.execute({"raw": "hello world"})
    print(f"Pipeline '{pipeline.name}' finished.")
    print(f"Final outputs: {outputs}")