在单元测试中模拟python类并验证实例

Question

I'm trying to unit test an SFTP helper class that makes some calls to the pysftp module. I want to mock the actual network calls from pysftp so there are no side effects, and just make sure that the class is correctly calling the underlying SFTP methods with the correct paramaters.

Here is a simple example of my code so far:

import pysftp
import unittest
import mock

class SFTPHelper(object):
    def __init__(self, host, username, password, files_dir):
        self.host = host
        self.username = username
        self.password = password
        self.files_dir = files_dir

    def list_files(self):
        with pysftp.Connection(
                self.host,
                username=self.username,
                password=self.password) as sftp:
            return sftp.listdir(self.files_dir)

class TestSFTPHelper(unittest.TestCase):
    @mock.patch('pysftp.Connection')
    def test_list_files(self, mock_connection):
        sftp_helper = SFTPHelper('somehost', 'someuser', 'somepassword', '/some/files/dir')
        sftp_helper.list_files()

        # this assertion passes
        mock_connection.assert_called_with(
            'somehost', password='somepassword', username='someuser')

        # this assertion does not pass
        mock_connection.listdir.assert_called_with('/some/files/dir')

The assertion error:

AssertionError: Expected call: listdir('/some/files/dir')
Not called

I assume it doesn't work because I need to assert that the function was called on the instance, but how do I get the instances of pysftp.Connection that was used in my method?

Answer 1

You could configure the mock to return a new mock object with __enter__ and __exit__ methods defined. For example:

@mock.patch.object(
    target=pysftp,
    attribute='Connection',
    autospec=True,
    return_value=mock.Mock(
        spec=pysftp.Connection,
        __enter__=lambda self: self,
        __exit__=lambda *args: None
    )
)
def test_list_files(self, mock_connection):
    # (contents of test case)

In addition, you may want to use:

mock_connection.return_value.listdir.assert_called_with('/some/files/dir')

instead of:

mock_connection.listdir.assert_called_with('/some/files/dir')

As a side note, you could also replace both uses of assert_called_with in your example with assert_called_once_with .

The end result:

$ python -m unittest test_sftp_helper.TestSFTPHelper.test_list_files
.
----------------------------------------------------------------------
Ran 1 test in 0.017s

OK