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.
Python code
60 linesimport 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
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
- Use `pathlib.Path.mkdir` with `parents=True` and `exist_ok=True` for nested directories
- 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
More from ML engineering pipelines
- Bayesian Optimization in Python: A Simplified Mock Implementation medium
- Build a Data Helper Class in Python for ML Pipelines easy
- Build a Mock Random Forest Classifier in Python easy
- Champion Challenger Deployment Mock in Python easy
- Compare Model A vs Model B Metrics in Python easy
- Create a Minimal Great Expectations Suite Mock in Python easy
Keep learning
Related tutorials and quizzes for this topic.