How to Mock MLflow Model Registration in Python

Build a lightweight in-memory mock of MLflow's MlflowClient to test model registration, versioning, and stage transitions without a tracking server.

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

mlflow mocking model-registry testing unit-test
Requires third-party packages — install first
pip install mlflow

Python code

42 lines
Python 3.9+ 
from mlflow.tracking import MlflowClient
from mlflow.entities import ModelVersion, Model


class MockMlflowClient:
    """Minimal mock of MlflowClient's model registration methods."""
    
    def __init__(self):
        self.registered_models = {}
        self.model_versions = {}
    
    def register_model(self, model_uri, name, stage=None):
        if name not in self.registered_models:
            self.registered_models[name] = Model(name=name)
        version = len(self.model_versions.get(name, [])) + 1
        mv = ModelVersion(name=name, version=version, source=model_uri, stage=stage or "None")
        self.model_versions.setdefault(name, []).append(mv)
        return mv
    
    def transition_model_version_stage(self, name, version, stage):
        mv = self.get_model_version(name, version)
        mv._stage = stage
        return mv
    
    def get_model_version(self, name, version):
        for mv in self.model_versions.get(name, []):
            if mv.version == version:
                return mv
        raise KeyError(f"Model version {version} not found for {name}")
    
    def list_model_versions(self, name):
        return self.model_versions.get(name, [])


if __name__ == "__main__":
    mock = MockMlflowClient()
    mv1 = mock.register_model("runs:/abc123/model", "iris_model")
    mv2 = mock.register_model("runs:/def456/model", "iris_model")
    print(f"Registered versions: {[v.version for v in mock.list_model_versions('iris_model')]}")
    
    mock.transition_model_version_stage("iris_model", 1, "Production")
    print(f"Stage of v1: {mock.get_model_version('iris_model', 1).stage}")

Output

stdout 
Registered versions: [1, 2]
Stage of v1: Production

How it works

This class wraps MlflowClient's key model registry methods — register, transition stage, get, and list — using plain dicts as storage. Each call mutates state in memory, so tests run fast with no external dependencies. ModelVersion is built from the real mlflow.entities class, so your code that handles version objects stays compatible. The mock is minimal by design: implement only the methods your pipeline calls, and you keep tests deterministic and isolated.

Common mistakes

  • Forgot to include the stage parameter in `__init__` of the mock, so it's missing from the returned object
  • Using a global state instead of a per-instance state, causing tests to leak data
  • Not raising `KeyError` for missing versions, which breaks `try/except` code paths
  • Mocking the whole class instead of just the methods you use, making setup overcomplicated

Variations

  1. Use `unittest.mock.patch` to replace `MlflowClient` with a `MagicMock` that records calls instead of the fake implementation
  2. Add `list_registered_models` or `delete_model_version` methods to the mock when your registry workflow needs cleanup

Real-world use cases

  • Unit-testing an ML pipeline that registers models but must not hit a real tracking server in CI.
  • Validating stage-transition logic (Staging → Production) before promoting a model in a deployment job.
  • Simulating multi-version registry behavior in integration tests that check which run artifact gets served.

Sponsored

Run locally

This sample needs third-party packages, so it cannot run in the browser IDE. Copy the code above, install the packages shown at the top, then run it in your own Python environment.

More from ML engineering pipelines

Keep learning

Related tutorials and quizzes for this topic.

Tutorials

Quizzes