Python单元测试：深度解析MLflow模型加载的Mocking策略

本文深入探讨了在python单元测试中，如何有效模拟mlflow模型加载（`mlflow.pyfunc.load_model`）这一常见挑战。当外部依赖在类初始化阶段被调用时，传统的`@patch`装饰器可能失效。文章通过分析问题根源，提出并演示了结合使用装饰器与`with patch`上下文管理器的解决方案，确保模型加载过程被可靠地模拟，从而实现对核心业务逻辑的独立测试。

问题描述与分析

在开发涉及MLflow模型加载的Python应用时，为了确保代码的健壮性和可测试性，我们通常需要对外部依赖（如MLflow的API调用）进行模拟（Mocking）。考虑以下一个类，它在初始化时会加载一个MLflow模型：

import mlflow

class MyClass:
   def __init__(self):
      # 假设 find_most_recent_model() 会返回一个有效的 run_id
      run_id = self.find_most_recent_model() 
      logged_model = f'mlruns/0/{run_id}/artifacts/pipeline'
      self.loaded_model = mlflow.pyfunc.load_model(logged_model)

   def find_most_recent_model(self):
       # 实际实现中会调用 MlflowClient.search_runs 等
       # 这里为了简化，假设其内部已处理并返回 run_id
       return 'mock_run_id' 

为了测试MyClass的初始化逻辑，我们需要模拟mlflow.MlflowClient.search_runs（如果find_most_recent_model依赖它）和mlflow.pyfunc.load_model，以避免实际的文件系统操作和MLflow服务器调用。最初的测试尝试可能如下：

from unittest.mock import Mock, patch
import pytest
import mlflow

# 假设MyClass在同一个模块或已正确导入
# from your_module import MyClass 

class Test:
    @patch('mlflow.MlflowClient.search_runs', Mock(return_value=[Mock(_info=Mock(run_id='mock_run_id'))]))
    @patch('mlflow.pyfunc.load_model', Mock()) # 尝试模拟 load_model
    def test_my_class_initialization(self):
        # 实例化 MyClass，期望所有MLflow调用都被模拟
        instance_test = MyClass() 
        # ... 后续断言 ...

在这种情况下，我们发现@patch('mlflow.MlflowClient.search_runs', ...)装饰器成功地模拟了find_most_recent_model内部的MLflow客户端调用，但@patch('mlflow.pyfunc.load_model', Mock())却未能生效。当MyClass实例化时，mlflow.pyfunc.load_model仍然尝试执行真实的逻辑，并因为路径'mlruns/0/mock_run_id/artifacts/pipeline'不存在而抛出OSError: No such file or directory错误。

这表明，对于在类初始化阶段（__init__方法）调用的外部函数，仅仅使用类或方法级别的@patch装饰器可能不足以确保模拟生效。这通常是由于Python的导入机制和patch的作用域问题导致的：如果被测试的类在模块加载时就已经导入了真实的mlflow模块，那么后续的@patch可能无法替换掉已经加载到MyClass作用域中的mlflow.pyfunc.load_model引用。

立即学习“Python免费学习笔记（深入）”；

度加剪辑

度加剪辑（原度咔剪辑），百度旗下AI创作工具

63

查看详情

解决方案：结合使用装饰器与上下文管理器

解决此类问题的最可靠方法之一是，对于那些在对象创建或关键操作期间必须被模拟的函数，使用unittest.mock.patch作为上下文管理器（with patch(...)）来明确限定其作用域。这确保了在代码块执行期间，目标函数确实被模拟了。

下面是修正后的测试代码：

from unittest.mock import Mock, patch
import pytest
import mlflow

# 假设MyClass在同一个模块或已正确导入
# from your_module import MyClass 

class MyClass:
   def __init__(self):
      run_id = self.find_most_recent_model()
      logged_model = f'mlruns/0/{run_id}/artifacts/pipeline'
      self.loaded_model = mlflow.pyfunc.load_model(logged_model)

   def find_most_recent_model(self):
       # 模拟 MlflowClient.search_runs 的调用
       # 实际代码中可能更复杂，这里简化
       return mlflow.MlflowClient().search_runs('0').pop()._info.run_id

class Test:
    # 装饰器用于模拟 MlflowClient.search_runs
    @patch('mlflow.MlflowClient.search_runs', Mock(return_value=[Mock(_info=Mock(run_id='mock_run_id'))]))
    def test_my_class_initialization_with_context_manager_patch(self):
        # 使用 with patch 作为上下文管理器，确保在 MyClass 实例化时 load_model 被模拟
        with patch('mlflow.pyfunc.load_model', return_value=Mock()) as mock_load_model:
            # 实例化 MyClass，此时 mlflow.pyfunc.load_model 将被 mock_load_model 替换
            instance_test = MyClass()

            # 验证 mock 是否被调用
            mock_load_model.assert_called_once()
            mock_load_model.assert_called_with('mlruns/0/mock_run_id/artifacts/pipeline')

            # 验证 loaded_model 是否是模拟对象
            assert instance_test.loaded_model is mock_load_model.return_value
            # ... 后续对 instance_test 的逻辑进行断言 ...

代码解析：

1. @patch('mlflow.MlflowClient.search_runs', ...): 这个装饰器继续有效，因为它模拟的是MyClass.find_most_recent_model内部可能使用的MlflowClient实例的search_runs方法。这个调用通常发生在__init__的早期，并且其目标路径是正确的。
2. with patch('mlflow.pyfunc.load_model', return_value=Mock()) as mock_load_model::
   • 上下文管理器: with patch语句创建了一个明确的模拟作用域。在with块内部，mlflow.pyfunc.load_model会被替换为我们提供的模拟对象。一旦离开with块，原始的load_model函数会自动恢复。
   • return_value=Mock(): 这一步至关重要。mlflow.pyfunc.load_model函数通常会返回一个模型对象，后续代码可能会对这个模型对象进行操作。通过设置return_value=Mock()，我们确保MyClass.loaded_model接收到一个可用的模拟对象，而不是None，从而避免后续属性访问错误。
   • as mock_load_model: 这允许我们获取对模拟对象的引用，以便在测试中进行断言，例如检查它是否被调用以及调用参数是否正确。
3. 实例化MyClass: 在with块内部实例化MyClass，确保在MyClass.__init__执行时，mlflow.pyfunc.load_model处于被模拟状态。

最佳实践与注意事项

1. 明确Patch目标路径: patch的目标路径必须是代码中实际查找和调用目标函数或对象的地方。如果MyClass是从my_module.py导入的，并且它内部调用了mlflow.pyfunc.load_model，那么正确的patch目标可能是my_module.mlflow.pyfunc.load_model，而不是直接mlflow.pyfunc.load_model。然而，在示例中，MyClass直接使用mlflow，因此mlflow.pyfunc.load_model是正确的。
2. 理解Patch作用域:
   • 装饰器@patch: 通常在整个测试函数或测试类生命周期内有效。对于在被测试函数内部调用的外部函数，通常效果良好。
   • 上下文管理器with patch: 提供了一个更精确和局部的模拟作用域。它在进入with块时应用模拟，在退出时恢复，非常适合处理在特定代码块（如对象初始化）中发生的调用。
3. 模拟返回值的重要性: 始终考虑被模拟函数或方法返回值的预期用途。如果被测试代码会使用返回对象，请确保mock的return_value是一个适当的Mock对象，或者是一个符合预期的真实值，以避免后续的AttributeError。
4. 避免过度模拟: 单元测试的目的是测试单个单元的逻辑，而不是其所有依赖项的内部工作原理。只模拟那些外部依赖，而不要模拟被测试单元自身的内部方法，除非这些内部方法本身是另一个单元。
5. 测试验证: 在使用patch后，务必添加断言来验证模拟对象是否按照预期被调用，例如使用mock_object.assert_called_once()、mock_object.assert_called_with(...)等。

总结

在Python单元测试中，当外部依赖（尤其是那些在对象初始化阶段被调用的函数）未能被@patch装饰器有效模拟时，结合使用@patch装饰器和with patch上下文管理器是一种强大而可靠的解决方案。with patch确保了在代码执行的关键时刻，目标函数被精确地替换为模拟对象，从而有效地隔离了被测试单元，使得测试更加专注于其自身的逻辑，提高了测试的稳定性和准确性。

以上就是Python单元测试：深度解析MLflow模型加载的Mocking策略的详细内容，更多请关注php中文网其它相关文章！