How to mock subsequent function calls in python?

Question

I'm new to testing and testing in python. I have a python class that looks like this :

File name : my_hive.py

from pyhive import hive

class Hive:
    def __init__(self, hive_ip):
        self.cursor = hive.connect(hive_ip).cursor()

    def execute(self, command):
        self.cursor.execute(command)

I want to mock these functions : pyhive.hive.connect, pyhive.Connection.cursor(used by my class as hive.connect(hive_ip).cursor()) and pyhive.Cursor.execute (used by my class as self.cursor.execute(command) in execute method).

I'm able to mock function call hive.connect and also I have been able to assert that it has been called with hive_ip given by me as follows.

import unittest
import mock

from my_hive import Hive

class TestHive(unittest.TestCase):
    @mock.patch('pyhive.hive.connect')
    def test_workflow(self, mock_connect):
        hive_ip = "localhost"
        processor = Hive(hive_ip)

        mock_connect.assert_called_with(hive_ip)

But how do I make sure that subsequent function calls like .cursor() and self.cursor.execute() have also been called? hive.connect(hive_ip) returns an instance of pyhive.hive.Connection, which has method called cursor

I have tried to add mocks like this :

import unittest
import mock

from hive_schema_processor import HiveSchemaProcessor

class TestHive(unittest.TestCase):
    @mock.patch('pyhive.hive.connect')
    @mock.patch('pyhive.hive.Connection.cursor')
    def test_workflow(self, mock_connect, mock_cursor):
        hive_ip = "localhost"
        processor = Hive(hive_ip)

        mock_connect.assert_called_with(hive_ip)
        mock_cursor.assert_called()

But the tests are failed with complain :

AssertionError: expected call not found.
Expected: cursor('localhost')
Actual: not called.

Answer 1

Your problem is that you have already mocked connect, so the subsequent calls on the result of connect will be made on the mock, not on the real object. To check that call, you have to make the check on the returned mock object instead:

class TestHive(unittest.TestCase):
    @mock.patch('pyhive.hive.connect')
    def test_workflow(self, mock_connect):
        hive_ip = "localhost"
        processor = Hive(hive_ip)

        mock_connect.assert_called_with(hive_ip)
        mock_cursor = mock_connect.return_value.cursor
        mock_cursor.assert_called()

Each call on a mock produces another mock object. mock_connect.return_value gives you the mock that is returned by calling mock_connect, and mock_connect.return_value.cursor contains another mock that will actually be called.