Mocking python class in unit test and verifying an instance

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?

1. Can you give a clearer problem description than *"doesn't work"*? 2. You're trying to mock a *context manager*, you should read up on that. — jonrsharpe, May 03 '16 at 21:57
@dnit13 just a typo in the stackoverflow question, I am doing that in the actual code. fixed — Spencer Wood, May 03 '16 at 22:09

score 9 · Accepted Answer · answered May 04 '16 at 05:45

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

Mocking python class in unit test and verifying an instance

1 Answers1

Linked