"Python：模擬上下文管理器"

Question

我不明白為什么我不能在這個例子中模擬 NamedTemporaryFile.name：

from mock import Mock, patch
import unittest
import tempfile

def myfunc():
    with tempfile.NamedTemporaryFile() as mytmp:
        return mytmp.name

class TestMock(unittest.TestCase):
    @patch('tempfile.NamedTemporaryFile')
    def test_cm(self, mock_tmp):
        mytmpname = 'abcde'
        mock_tmp.__enter__.return_value.name = mytmpname
        self.assertEqual(myfunc(), mytmpname)

Answer 1

您設置了錯誤的模擬： mock_tmp不是上下文管理器，而是返回一個上下文管理器。 將您的設置行替換為：

mock_tmp.return_value.__enter__.return_value.name = mytmpname

並且您的測試將起作用。

Answer 2

使用 pytest 和模擬裝置擴展 Peter K 的答案。

def myfunc():
    with tempfile.NamedTemporaryFile(prefix='fileprefix') as fh:
        return fh.name


def test_myfunc(mocker):
    mocker.patch('tempfile.NamedTemporaryFile').return_value.__enter__.return_value.name = 'tempfilename'
    assert myfunc() == 'tempfilename'

Answer 3

這是pytest和mocker fixture的替代方案，這也是一種常見做法：

def test_myfunc(mocker):
    mock_tempfile = mocker.MagicMock(name='tempfile')
    mocker.patch(__name__ + '.tempfile', new=mock_tempfile)
    mytmpname = 'abcde'
    mock_tempfile.NamedTemporaryFile.return_value.__enter__.return_value.name = mytmpname
    assert myfunc() == mytmpname

Answer 4

為了擴展納撒尼爾的答案，這個代碼塊

with tempfile.NamedTemporaryFile() as mytmp:
    return mytmp.name

有效地做三件事

# Firstly, it calls NamedTemporaryFile, to create a new instance of the class.
context_manager = tempfile.NamedTemporaryFile()  

# Secondly, it calls __enter__ on the context manager instance.
mytmp = context_manager.__enter__()  

# Thirdly, we are now "inside" the context and can do some work. 
return mytmp.name

當您將tempfile.NamedTemporaryFile替換為Mock或MagicMock的實例時

# This first line, below, will call mock_tmp().
# Therefore we need to set the return_value with
# mock_tmp.return_value
context_manager = mock_tmp()

# This will call mock_tmp.return_value.__enter__() so we need to set 
# mock_tmp.return_value.__enter__.return_value
mytmp = context_manager.__enter__()

# This will access mock_tmp.return_value.__enter__.return_value.name
return mytmp.name

Answer 5

我將 hmobrienv 的回答擴展到一個小型工作程序

import tempfile
import pytest


def myfunc():
    with tempfile.NamedTemporaryFile(prefix="fileprefix") as fh:
        return fh.name


def test_myfunc(mocker):
    mocker.patch("tempfile.NamedTemporaryFile").return_value.__enter__.return_value.name = "tempfilename"
    assert myfunc() == "tempfilename"


if __name__ == "__main__":
    pytest.main(args=[__file__])

Answer 6

另一種可能性是使用工廠創建實現上下文管理器接口的對象：

import unittest
import unittest.mock
import tempfile


def myfunc():
    with tempfile.NamedTemporaryFile() as mytmp:
        return mytmp.name


def mock_named_temporary_file(tmpname):
    class MockNamedTemporaryFile(object):
        def __init__(self, *args, **kwargs):
            self.name = tmpname

        def __enter__(self):
            return self

        def __exit__(self, type, value, traceback):
            pass

    return MockNamedTemporaryFile()


class TestMock(unittest.TestCase):
    @unittest.mock.patch("tempfile.NamedTemporaryFile")
    def test_cm(self, mock_tmp):
        mytmpname = "abcde"
        mock_tmp.return_value = mock_named_temporary_file(mytmpname)
        self.assertEqual(myfunc(), mytmpname)