How to mock an artifact store with local paths in Python for ML pipelines

Create a temporary local artifact store with dummy files and metadata to test ML pipeline code without real storage.

Medium Python 3.9+ Aug 9, 2026 ML engineering pipelines 13 views 0 copies

Python code

60 lines
Python 3.9+
import tempfile
from pathlib import Path
import json


def create_artifact_store_mock(base_path: Path = None):
    """Create a local artifact store mock directory structure."""
    if base_path is None:
        base_path = Path(tempfile.mkdtemp())

    store_layout = {
        "artifacts": [
            {"name": "model.pkl", "size": 1024},
            {"name": "metrics.json", "size": 512},
            {"name": "training_data.csv", "size": 2048}
        ],
        "metadata": {
            "project": "demo",
            "version": "1.0.0"
        }
    }

    artifact_dir = base_path / "artifacts"
    artifact_dir.mkdir(parents=True, exist_ok=True)
    metadata_dir = base_path / "metadata"
    metadata_dir.mkdir(exist_ok=True)

    for artifact in store_layout["artifacts"]:
        artifact_path = artifact_dir / artifact["name"]
        content = b"0" * artifact["size"]
        artifact_path.write_bytes(content)

    metadata_path = metadata_dir / "store_config.json"
    metadata_path.write_text(json.dumps(store_layout["metadata"], indent=2))

    return base_path


def inspect_mock_store(store_path: Path):
    """Inspect and print details of the mock artifact store."""
    artifacts_dir = store_path / "artifacts"
    metadata_file = store_path / "metadata" / "store_config.json"

    if not artifacts_dir.is_dir() or not metadata_file.is_file():
        raise ValueError("Invalid mock store structure")

    artifact_files = sorted(artifacts_dir.iterdir())
    metadata = json.loads(metadata_file.read_text())

    print(f"Store root: {store_path}")
    print(f"Artifacts ({len(artifact_files)}):")
    for artifact in artifact_files:
        print(f"  - {artifact.name} ({artifact.stat().st_size} bytes)")

    print(f"Metadata: {json.dumps(metadata)}")


if __name__ == "__main__":
    store_path = create_artifact_store_mock()
    inspect_mock_store(store_path)

Output

stdout
Store root: /tmp/tmpabcdefgh
Artifacts (3):
  - metrics.json (512 bytes)
  - model.pkl (1024 bytes)
  - training_data.csv (2048 bytes)
Metadata: {"project": "demo", "version": "1.0.0"}

How it works

The function uses tempfile.mkdtemp() to create a unique temporary directory, ensuring isolated test runs. Path.write_bytes and write_text provide clean file I/O without manual open/close management. The metadata JSON is serialized with json.dumps for human-readable config. sorted() on directory iteration gives deterministic ordering of artifacts across platforms. This pattern lets you swap the mock for a real cloud artifact store later without changing downstream code.

Common mistakes

  • Forgetting to clean up temporary directories after tests
  • Assuming file sizes must be realistic — using `b'0' * size` is fine for most mock scenarios
  • Using `os.makedirs` without `exist_ok=True` causes crashes on re-runs
  • Not sorting `iterdir()` results, leading to non-deterministic output

Variations

  1. Use `pathlib.Path.mkdir` with `parents=True` and `exist_ok=True` for nested directories
  2. Replace `tempfile.mkdtemp` with `pytest tmp_path` fixture for automatic cleanup

Real-world use cases

  • Testing ML training code locally when the production artifact store (S3/GCS) is unavailable in CI.
  • Simulating model registry paths for unit tests that validate artifact metadata before deployment.
  • Building integration tests for pipeline steps that read model files and metrics without provisioning real cloud storage.

Sponsored

Run this sample

Open the browser IDE to tweak the example and see results without installing anything.

Open editor

More from ML engineering pipelines

Related tutorials and quizzes for this topic.