Mocking method that calls other methods

Question

I have a class

class X:

  def __init__(self, db):
    self.db = db

  def get_data_from_friend(self):
    return None

  def get_data_from_db(self):
    return self.db.get_my_db_data()

  def get_data(self):
    if data := self.get_data_from_friend():
      return data

    return self.get_data_from_db()

And I'm trying to test the get_data method, to validate that the calls inside of it are executed.

I have a test like this

def test_get_data(self):
  mock = create_autospec(X)
  mock.get_data()
  mock.get_data.assert_called_once() # <-- works
  mock.get_data_from_friend.assert_called_once() # <-- assertionError, not called

what am I missing here?

Answer 1

Your X class is correctly defined and accepts dependency injection so you can do something like this

from unittest.mock import MagicMock

class X:

  def __init__(self, db):
    self.db = db

  def get_data_from_friend(self):
    return None

  def get_data_from_db(self):
    return self.db.get_my_db_data()

  def get_data(self):
    if data := self.get_data_from_friend():
      return data

    return self.get_data_from_db()

MOCK_DATA = "mock data"
def test_get_data():
    db_mock = MagicMock()
    db_mock.get_my_db_data = MagicMock(return_value=MOCK_DATA)
    assert X(db_mock).get_data() == MOCK_DATA