Python中如何测试SFTP连接及文件传输方法

你想得没错——用@mock.patch()来伪造SFTP连接正是解决这个问题的关键，这样就能避免实际建立连接和传输文件，只验证代码的逻辑行为。下面我给你详细的测试方案和代码示例：

我们需要通过unittest.mock（或pytest-mock）模拟pysftp.Connection类，让它在测试时不真实初始化SFTP连接，而是返回一个模拟对象。重点验证以下几点：

• Connection是否使用正确的参数（host、username、password）被调用
• put方法是否被触发，且传入了正确的文件路径
• 异常场景下方法是否返回None并正确记录日志

假设你的方法属于FileTransfer类，下面分两种常用测试框架给出实现：

1. 使用unittest.mock（标准库） #

import unittest
from unittest.mock import patch, MagicMock
from your_module import FileTransfer  # 替换为你的实际模块路径

class TestFileTransfer(unittest.TestCase):
    def test_transfer_file_success(self):
        # 初始化测试实例和模拟文件路径
        transfer = FileTransfer()
        test_file_path = "/tmp/test_upload.txt"

        # 模拟pysftp.Connection类及其上下文管理器
        with patch('your_module.pysftp.Connection') as mock_conn:
            # 配置mock的sftp对象（即with语句中as后的变量）
            mock_sftp = MagicMock()
            mock_conn.return_value.__enter__.return_value = mock_sftp

            # 调用待测试方法
            result = transfer._transfer_file(test_file_path)

            # 验证关键行为
            mock_conn.assert_called_once_with('host', username=username, password=password)
            mock_sftp.put.assert_called_once_with(test_file_path)
            self.assertEqual(result, True)

    def test_transfer_file_failure(self):
        transfer = FileTransfer()
        test_file_path = "/tmp/test_upload.txt"

        with patch('your_module.pysftp.Connection') as mock_conn:
            # 模拟连接时抛出异常
            mock_conn.return_value.__enter__.side_effect = Exception("SFTP连接超时")

            # 调用方法
            result = transfer._transfer_file(test_file_path)

            # 验证异常处理逻辑
            self.assertIsNone(result)
            # 若有日志记录，可额外验证日志是否被触发
            # mock_log = patch('your_module.logging.error').start()
            # mock_log.assert_called_once()

2. 使用pytest-mock（更简洁的pytest方案） #

如果你用pytest框架，配合pytest-mock插件可以简化代码：

from your_module import FileTransfer

def test_transfer_file_success(mocker):
    transfer = FileTransfer()
    test_file_path = "/tmp/test_upload.txt"

    # 模拟sftp连接对象
    mock_sftp = mocker.Mock()
    mock_conn = mocker.patch('your_module.pysftp.Connection')
    mock_conn.return_value.__enter__.return_value = mock_sftp

    # 执行测试
    result = transfer._transfer_file(test_file_path)

    # 验证行为
    mock_conn.assert_called_once_with('host', username=username, password=password)
    mock_sftp.put.assert_called_once_with(test_file_path)
    assert result is True

def test_transfer_file_failure(mocker):
    transfer = FileTransfer()
    test_file_path = "/tmp/test_upload.txt"

    # 模拟连接异常
    mock_conn = mocker.patch('your_module.pysftp.Connection')
    mock_conn.return_value.__enter__.side_effect = Exception("权限不足")

    result = transfer._transfer_file(test_file_path)

    assert result is None

• Mock路径要准确：必须patch你代码中实际导入pysftp.Connection的路径。比如如果你的代码是from pysftp import Connection，就要patchyour_module.Connection；如果是import pysftp后用pysftp.Connection，则patchyour_module.pysftp.Connection。这是mock最容易踩坑的地方。
• 上下文管理器的模拟：因为使用了with ... as sftp语法，必须模拟__enter__方法的返回值，这个返回值就是代码中用到的sftp对象。
• 覆盖异常场景：除了成功路径，一定要测试连接失败、文件不存在等异常情况，确保方法的错误处理逻辑正常工作。

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