如何在Nornir 3.3.0中使用unittest.mock模拟netmiko_send_command？

问题描述 #

我正在使用Python 3.8搭配Nornir（3.3.0）自动化框架，希望模拟对设备的SSH访问，无需真实或虚拟网络设备在线即可完成测试。请问如何使用unittest.mock中的patch、Mock/MagicMock来模拟netmiko_send_command（与设备的SSH交互）？

我的任务函数 #

# dbb_automation/tasks.py

from nornir.core.task import Task, Result
from nornir_netmiko.tasks import netmiko_send_command

def get_interfaces_with_ip(task: Task):
    log.debug(f"{task.name}: Getting result on host {task.host}")
    result: MultiResult = task.run(name="show ip int br | e unass", task=netmiko_send_command,
                      command_string="show ip int br | e unass")
    
    content_str = result[0].result

    task.run(
        task=write_file,
        filename=f"outputs/{task.host}-{purpose}.{ending}",
        content=content_str
    )
    
    return Result(
        host=task.host,
        result=f"{task.host.name} got ip result"
    )

待完善的测试用例 #

# tests/test_tasks.py

from dbb_automation.tasks import get_interfaces_with_ip
from nornir import InitNornir
from nornir.core.filter import F
from tests.settings import *

def test_get_interfaces_with_ip():
    # [x] init nornir with fake host
    # [ ] patch/mock netmiko_send_command
    # [ ] check file contents with patched return string of netmiko_send_command

    nr = InitNornir(
        core={
            "raise_on_error": True
        },
        runner={
            "plugin": "threaded",
            "options": {
                "num_workers": 1,
            }
        },
        inventory={
            "plugin": "SimpleInventory",
            "options": {
                "host_file": DNAC_HOSTS_YAML,
                "group_file": DNAC_GROUPS_YAML,
                "defaults_file": DNAC_DEFAULT_YAML
            }
        },
        logging={
            "log_file": "logs/nornir.log"
        }
    )
    result = nr.filter(F(has_parent_group="Borders")).run(name="get_interfaces_with_ip", task=get_interfaces_with_ip)
    
    # todo: test code

    assert False

解决方案 #

核心思路 #

• 用unittest.mock.patch替换任务函数中调用的netmiko_send_command，模拟它的返回值
• 构造符合Nornir MultiResult结构的Mock对象，让任务能正确获取result[0].result
• 可选：Mock write_file任务，避免生成真实文件，同时验证它是否被正确调用

完善后的测试用例 #

# tests/test_tasks.py

from unittest.mock import patch, MagicMock
from dbb_automation.tasks import get_interfaces_with_ip
from nornir import InitNornir
from nornir.core.filter import F
from nornir.core.task import Result
from tests.settings import *

def test_get_interfaces_with_ip():
    # 1. 预设模拟的命令输出
    mock_command_output = """Interface              IP-Address      OK? Method Status                Protocol
GigabitEthernet0/0     192.168.1.1     YES NVRAM  up                    up      
GigabitEthernet0/1     10.0.0.1        YES NVRAM  up                    up"""

    # 2. 构造Mock的netmiko_send_command返回逻辑
    def mock_netmiko_send_command(task, **kwargs):
        # 返回模拟的Result对象，模拟netmiko执行命令的结果
        return Result(host=task.host, result=mock_command_output)

    # 3. 初始化Nornir（保留原有逻辑）
    nr = InitNornir(
        core={
            "raise_on_error": True
        },
        runner={
            "plugin": "threaded",
            "options": {
                "num_workers": 1,
            }
        },
        inventory={
            "plugin": "SimpleInventory",
            "options": {
                "host_file": DNAC_HOSTS_YAML,
                "group_file": DNAC_GROUPS_YAML,
                "defaults_file": DNAC_DEFAULT_YAML
            }
        },
        logging={
            "log_file": "logs/nornir.log"
        }
    )
    target_hosts = nr.filter(F(has_parent_group="Borders"))

    # 4. 使用patch替换目标函数
    with patch("dbb_automation.tasks.netmiko_send_command", side_effect=mock_netmiko_send_command) as mock_send_cmd, \
         patch("dbb_automation.tasks.write_file") as mock_write:
        
        # 运行任务
        result = target_hosts.run(name="get_interfaces_with_ip", task=get_interfaces_with_ip)

        # 5. 验证netmiko_send_command的调用情况
        assert mock_send_cmd.called
        # 检查命令参数是否正确
        call_args = mock_send_cmd.call_args
        assert call_args[1]["command_string"] == "show ip int br | e unass"

        # 6. 验证write_file的调用情况
        assert mock_write.called
        # 检查写入内容是否匹配模拟输出
        write_call_args = mock_write.call_args
        assert write_call_args[1]["content"] == mock_command_output
        # 验证文件名包含目标主机名
        for host_name in target_hosts.inventory.hosts:
            assert host_name in write_call_args[1]["filename"]

        # 7. 验证任务返回结果符合预期
        for host_name, host_result in result.items():
            assert host_result.result == f"{host_name} got ip result"
            assert not host_result.failed

关键细节说明 #

• patch路径正确性：必须patch任务函数中实际导入netmiko_send_command的路径（即dbb_automation.tasks.netmiko_send_command），而非nornir_netmiko.tasks.netmiko_send_command，因为Python的patch基于模块命名空间生效。
• 模拟Result结构：task.run调用任务会自动返回MultiResult，所以我们的Mock函数直接返回Result即可，任务里的result[0].result能正确获取模拟输出。
• 调用验证：通过mock_send_cmd.called和call_args可以确认模拟函数是否被触发，以及参数是否符合预期，确保任务逻辑正确。
• Mock写入操作：如果无需测试真实文件写入，Mockwrite_file可以避免生成测试文件，同时验证任务是否正确触发了写入流程。

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