如何用unittest mock优雅捕获with代码块内的方法调用？

import pysftp
from unittest import mock

remote_file_name = 'remote_file_name'
local_path = 'local/path/to/file'

def open_sftp_connection(hostname):
    return pysftp.Connection(hostname, username='usr', password='pswd')

def sftp_upload_try(local_path, dest_file_name):
    sftp_conn = None
    try:
        sftp_conn = open_sftp_connection('sftp-host')
        sftp_conn.put(local_path, remotepath=dest_file_name)
    finally:
        if sftp_conn:
            sftp_conn.close()

def test_sftp_try():
    conn_mock = mock.MagicMock()
    with mock.patch('test_sftp.open_sftp_connection', return_value=conn_mock):
        sftp_upload_try(local_path, remote_file_name)
        print(conn_mock.mock_calls)
        assert 1 == conn_mock.put.call_count
        assert remote_file_name == conn_mock.put.call_args.kwargs['remotepath']
        assert 1 == conn_mock.close.call_count

print(conn_mock.mock_calls)的输出：

[call.put('local/path/to/file', remotepath='remote_file_name'),
 call.__bool__(),
 call.close()]

def sftp_upload_with(local_path, dest_file_name):
    with open_sftp_connection('sftp-host') as sftp_conn:
        sftp_conn.put(f'${local_path}', remotepath=dest_file_name)

def test_sftp_with():
    conn_mock = mock.MagicMock()
    with mock.patch('test_sftp.open_sftp_connection', return_value=conn_mock):
        sftp_upload_with(local_path, remote_file_name)
        print(conn_mock.mock_calls)
        put_call = next(c for c in conn_mock.mock_calls if '__enter__().put' == c[0])
        assert put_call
        assert remote_file_name == put_call.kwargs['remotepath']

print(conn_mock.mock_calls)的输出：

[call.__enter__(),
 call.__enter__().put('$local/path/to/file', remotepath='remote_file_name'),
 call.__exit__(None, None, None)]

问题 #

使用with语句时，有没有更优雅的方式断言方法调用？还是只能用类似'__enter__().put' == c[0]这种不够简洁的方式？

解答 #

有两种更简洁优雅的方式，不需要遍历mock_calls：

1. 直接链式断言__enter__()返回对象的方法
   with语句会调用连接对象的__enter__()方法，并将返回值绑定到as后的变量。你可以直接通过conn_mock.__enter__().put访问该方法的mock对象，直接做断言：

def test_sftp_with():
    conn_mock = mock.MagicMock()
    with mock.patch('test_sftp.open_sftp_connection', return_value=conn_mock):
        sftp_upload_with(local_path, remote_file_name)
        # 断言put方法被调用一次
        conn_mock.__enter__().put.assert_called_once()
        # 断言put方法的参数完全匹配
        conn_mock.__enter__().put.assert_called_with(
            f'${local_path}', 
            remotepath=remote_file_name
        )

2. 提前提取__enter__()的返回mock
   如果需要多次断言，先把__enter__()返回的mock对象提取出来，代码会更清晰：

def test_sftp_with():
    conn_mock = mock.MagicMock()
    # 提取with语句实际使用的session mock
    sftp_session_mock = conn_mock.__enter__()
    with mock.patch('test_sftp.open_sftp_connection', return_value=conn_mock):
        sftp_upload_with(local_path, remote_file_name)
        sftp_session_mock.put.assert_called_once()
        assert remote_file_name == sftp_session_mock.put.call_args.kwargs['remotepath']
        # 还可以顺便断言__exit__方法的调用
        conn_mock.__exit__.assert_called_once_with(None, None, None)

这两种方式都利用了MagicMock的链式调用特性，完全避免了遍历调用记录的繁琐，代码更直观易读，也符合unittest.mock的设计逻辑。

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