python - 如何使用 pytest 测试异常

Question

如何在不调用实际 ftp 服务器的情况下测试我的 _fetch_files 的异常分支？下面是我的实现代码和正在测试异常分支的当前测试。

def _fetch_files(ftp_server:str, ftp_dir:str, file_name:str, dir_path:pathlib.Path) -> None:
    ''' logs into noaa's ftp server and downloads to memory `.gz` files for a given year, for a given file_name
        Args:
            ftp_server: string of ftp server 
            ftp_dir: dir_path containing `.gz` files 
            file_name: weather station by `.gz` file 
            dir_path: the dir path to which the files will be saved to 
        Returns:
            None
    '''
    with ftplib.FTP(ftp_server, timeout=3.0) as ftp:
        ftp.login()
        ftp.cwd(ftp_dir)
        make_raw_dir(dir_path)
        try:
            with open(file_name, 'wb') as fp:
                ftp.retrbinary(f'RETR {file_name}', fp.write)
                logger.info(f'writing file: {file_name}')
        except ftplib.error_perm:
                logger.error(f'{file_name} not found')
        ftp.quit()
def test_fetch_files_exception(tmp_path): 
    tmp_dir_path = tmp_path / 'sub'
    tmp_dir_path.mkdir()
    with mock.patch('module_three.utils_IO_bound._fetch_files', side_effect=Exception) as mock_req:
        with pytest.raises(Exception):
            assert _fetch_files('ftp.ncdc.noaa.gov', 'pub/data/noaa/2016', '123.gz', tmp_dir_path) == '123.gz not found'

标签： pythonexceptionpytest

Answer 1

由于我们需要模拟的块是上下文管理器（此处为ftplib.FTP），因此我们必须控制以下内容以将流程引导到成功场景或异常场景。

contextmanager.__enter__()

...此方法返回的值绑定到使用此上下文管理器as的语句子句中的标识符...with

contextmanager.__exit__(exc_type, exc_val, exc_tb)

...从该方法返回一个true值将导致该with 语句抑制异常并继续执行该语句之后的with语句。否则，此方法完成执行后异常继续传播。...

在这里，我们将模拟ftplib.FTP并ftplib.FTP.retrbinary控制它是通过成功场景还是异常场景。

test_ftp.py

import ftplib
from unittest import mock

import pytest


# Simplified version to focus on the test logic of mocking FTP
def _fetch_files(ftp_server: str, ftp_dir: str, file_name: str) -> None:
    with ftplib.FTP(ftp_server, timeout=3.0) as ftp:
        print(f"FTP object {type(ftp)} {ftp}")
        ftp.login()
        ftp.cwd(ftp_dir)
        try:
            with open(file_name, 'wb') as fp:
                ftp.retrbinary(f'RETR {file_name}', fp.write)
                print(f'writing file: {file_name}')
        except ftplib.error_perm:
            print(f'{file_name} not found')
            # Raise an error so that we can see if the failure actually happened
            raise FileNotFoundError(f'{file_name} not found')
        ftp.quit()


@mock.patch("ftplib.FTP")  # Mock FTP so that we wouldn't access the real world files
def test_fetch_files_success(mock_ftp):
    # With all FTP operations mocked to run ok, the flow should go through the success scenario.
    _fetch_files('ftp.ncdc.noaa.gov', 'pub/data/noaa/2016', '123.gz')


@mock.patch("ftplib.FTP")  # Mock FTP so that we wouldn't access the real world files
def test_fetch_files_exception(mock_ftp):
    """
    Mock FTP:retrbinary to raise an exception. This should go through the exception scenario.
    1. mock_ftp.return_value
        The FTP object, here being the <ftplib.FTP(ftp_server, timeout=3.0)>
    2. .__enter__.return_value
        The object to be bound in the <as> clause, here being the <ftp> variable in <with ftplib.FTP(ftp_server, timeout=3.0) as ftp:>
    3. .retrbinary.side_effect
        The behavior if the bound object <ftp> is used to call <ftp.retrbinary(...)>, which here was configured to raise the exception <ftplib.error_perm>
    """
    mock_ftp.return_value.__enter__.return_value.retrbinary.side_effect = ftplib.error_perm

    with pytest.raises(FileNotFoundError) as error:
        _fetch_files('ftp.ncdc.noaa.gov', 'pub/data/noaa/2016', '123.gz')

    assert str(error.value) == '123.gz not found'

日志

$ pytest -rP
================================================================================================= PASSES ==================================================================================================
________________________________________________________________________________________ test_fetch_files_success _________________________________________________________________________________________
------------------------------------------------------------------------------------------ Captured stdout call -------------------------------------------------------------------------------------------
FTP object <class 'unittest.mock.MagicMock'> <MagicMock name='FTP().__enter__()' id='140376736032224'>
writing file: 123.gz
_______________________________________________________________________________________ test_fetch_files_exception ________________________________________________________________________________________
------------------------------------------------------------------------------------------ Captured stdout call -------------------------------------------------------------------------------------------
FTP object <class 'unittest.mock.MagicMock'> <MagicMock name='FTP().__enter__()' id='140376735714608'>
123.gz not found
============================================================================================ 2 passed in 0.03s ============================================================================================

参考：

• Python：模拟上下文管理器
• https://moagstar.github.io/posts/a-better-way-to-mock-context-managers-in-python/