How to Mock MLflow log_params and log_metrics in Python

Use unittest.mock to patch MLflow's log_param and log_metric, run the training function, and verify logging calls without touching a real tracking server.

Easy Python 3.9+ Aug 9, 2026 ML engineering pipelines 15 views 0 copies

mlflow mock testing unittest
Requires third-party packages — install first
pip install mlflow

Python code

21 lines
Python 3.9+ 
from unittest.mock import Mock, patch
import mlflow


def train_model():
    mlflow.log_param("learning_rate", 0.01)
    mlflow.log_param("epochs", 10)
    mlflow.log_metric("accuracy", 0.95)
    mlflow.log_metric("loss", 0.05)
    return "Training completed"


if __name__ == "__main__":
    with patch("mlflow.log_param") as mock_log_param, patch("mlflow.log_metric") as mock_log_metric:
        result = train_model()
        
        print(result)
        print(f"Params logged: {mock_log_param.call_count}")
        print(f"Metrics logged: {mock_log_metric.call_count}")
        print(f"First param call: {mock_log_param.call_args_list[0]}")
        print(f"First metric call: {mock_log_metric.call_args_list[0]}")

Output

stdout 
Training completed
Params logged: 2
Metrics logged: 2
First param call: call('learning_rate', 0.01)
First metric call: call('accuracy', 0.95)

How it works

unittest.mock.patch replaces the mlflow.log_param and mlflow.log_metric functions with Mock objects for the duration of the with block. This lets you run model training code without an active MLflow tracking URI, so it works in CI or on a laptop with no server configured. The mock records every call, and call_count and call_args_list expose how many times each function ran and with which arguments. After the with block exits, the originals are restored, so real logging resumes elsewhere. This pattern is ideal for unit tests of training scripts where you only care about the side effect of logging, not the actual MLflow backend.

Common mistakes

  • Forgetting that patch needs the exact import path used in the code (here 'mlflow.log_param', not 'log_param').
  • Asserting on mock.call_args without checking call_count first, leading to IndexError on empty lists.
  • Running training in production with mocks still active because they forgot the `with` block scope.

Variations

  1. Use `patch('mlflow.log_params', autospec=True)` to also mock the batch logging API for dict-style inputs.
  2. Add assertions like `mock_log_param.assert_any_call('epochs', 10)` to verify specific arguments were passed.

Real-world use cases

  • Writing pytest suites for training pipelines that must not hit a remote MLflow tracking server in CI.
  • Verifying that a feature-engineering script calls the correct logging functions with expected hyperparameters.
  • Testing model evaluation code in notebooks or jobs where MLflow is not installed or configured.

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