Python Mocking Netmiko

Question

I am building a netmiko wrapper for a project and I'm trying to write some simple unit tests to verify functionality. However I am having trouble mocking the netmiko library. I was able to mock the ConnectHandler object creation pretty easily, but verifying methods of the handler were called is not working as expected.

In the code below you'll see my simple wrapper and unit tests. The problem lies with test_disconnect. I am trying to verify the that disconnect method is called, but for some reason mock is reporting that it was not called. I'm sure this has something to do with the way I am mocking netmiko, but I am unsure what I'm doing wrong.

MockSandbox.py

from netmiko import ConnectHandler, NetMikoTimeoutException


class SshLib(object):
    def __init__(self, ip, user, passwd, port=22):
        self.ip = ip
        self.user = user
        self.passwd = passwd
        self.port = port
        self.connection = None

    def login(self):
        args = {"device_type": "cisco_ios",
                "ip": self.ip,
                "username": self.user,
                "password": self.passwd,
                "port": self.port}

        try:
            self.connection = ConnectHandler(**args)
            return True
        except NetMikoTimeoutException:
            return False

    def disconnect(self):
        self.connection.disconnect()

MockSandboxTests.py

import mock
import unittest
from MockSandbox import SshLib


class SshLibTests(unittest.TestCase):
    @mock.patch("MockSandbox.ConnectHandler")
    def test_connect(self, mock_ssh):
        ssh = SshLib("1.1.1.1", "user", "password")
        return_val = ssh.login()

        args = {"device_type": "cisco_ios",
                "ip": "1.1.1.1",
                "username": "user",
                "password": "password",
                "port": 22}

        mock_ssh.assert_called_with(**args)
        self.assertTrue(return_val)

    @mock.patch("MockSandbox.ConnectHandler")
    def test_disconnect(self, mock_ssh):
        ssh = SshLib("1.1.1.1", "user", "password")
        ssh.login()
        ssh.disconnect()
        mock_ssh.disconnect.assert_called()


if __name__ == '__main__':
    unittest.main()

Test Output

Failure
Traceback (most recent call last):
  File "C:\Python27\lib\unittest\case.py", line 329, in run
    testMethod()
  File "C:\Python27\lib\site-packages\mock\mock.py", line 1305, in patched
    return func(*args, **keywargs)
  File "C:\Users\user\Python\MockSandboxTests.py", line 26, in test_disconnect
    mock_ssh.disconnect.assert_called()
  File "C:\Python27\lib\site-packages\mock\mock.py", line 906, in assert_called
    raise AssertionError(msg)
AssertionError: Expected 'disconnect' to have been called.



Ran 2 tests in 0.020s

FAILED (failures=1)

Answer 1

mock.patch("MockSandbox.ConnectHandler") returns you a mock for the class, so mock_ssh.disconnect is a class method. You want to check the instance method instead. Access the mock instance via mock_ssh.return_value:

mock_ssh.return_value.disconnect.assert_called()