如何用pytest断言方法中的变量被赋值为正确的类实例？

你不需要直接获取未返回的registrar变量，只需通过mock Registrar类的构造函数，验证对应类是否被实例化，就能证明变量赋值逻辑正确。以下是具体实现方案：

核心思路 #

原函数中registrar是通过Registrar1()或Registrar2()实例化得到的，只要验证对应类的构造函数被调用（或未被调用），就能确认赋值逻辑是否符合预期。同时需要mockis_available方法返回固定值，避免原方法的实际逻辑干扰测试。

测试代码示例 #

假设目标函数和Registrar类在my_module.py中，使用pytest-mock（需提前安装：pip install pytest-mock）编写测试：

from my_module import domain_is_available, Registrar1, Registrar2

def test_jp_domain_uses_registrar1(mocker):
    # Mock Registrar1的构造函数及is_available方法
    mock_registrar1 = mocker.patch("my_module.Registrar1")
    mock_registrar1.return_value.is_available.return_value = True
    
    # Mock Registrar2的构造函数及is_available方法
    mock_registrar2 = mocker.patch("my_module.Registrar2")
    mock_registrar2.return_value.is_available.return_value = False
    
    # 执行目标函数
    domain_is_available("test.jp")
    
    # 断言Registrar1被正确实例化，Registrar2未被调用
    mock_registrar1.assert_called_once_with(ip_rotate=False)
    mock_registrar2.assert_not_called()

def test_non_jp_domain_uses_registrar2(mocker):
    mock_registrar1 = mocker.patch("my_module.Registrar1")
    mock_registrar2 = mocker.patch("my_module.Registrar2")
    mock_registrar2.return_value.is_available.return_value = False
    
    domain_is_available("test.com")
    
    mock_registrar2.assert_called_once_with(ip_rotate=False)
    mock_registrar1.assert_not_called()

# 测试带ip_rotate参数的场景
def test_jp_domain_with_ip_rotate(mocker):
    mock_registrar1 = mocker.patch("my_module.Registrar1")
    
    domain_is_available("test.jp", ip_rotate=True)
    
    mock_registrar1.assert_called_once_with(ip_rotate=True)

关键说明 #

1. mocker.patch会替换目标模块中的Registrar类为mock对象，构造函数的调用记录会被保存。
2. assert_called_once_with用于验证构造函数是否以正确的参数被调用，确保ip_rotate参数传递正确。
3. 无需关心is_available的实际返回值，只需mock一个固定布尔值让函数正常执行即可。

内容的提问来源于stack exchange，提问作者Alvaro