How can I run a function (to get side effects) when a python Mock is called?

Question

I am mocking (using python Mock) a function where I want to return a list of values, but in some of the items in the list I want a side effect to also occur (at the point where the mocked function is called). How is this most easily done? I'm trying something like this:

import mock
import socket

def oddConnect():
  result = mock.MagicMock()  # this is where the return value would go
  raise socket.error  # I want it assigned but also this raised

socket.create_connection = mock(spec=socket.create_connection,
  side_effect=[oddConnect, oddConnect, mock.MagicMock(),])
# what I want: call my function twice, and on the third time return normally
# what I get: two function objects returned and then the normal return

for _ in xrange(3):
  result = None
  try:
    # this is the context in which I want the oddConnect function call
    # to be called (not above when creating the list)
    result = socket.create_connection()
  except socket.error:
    if result is not None:
      # I should get here twice
      result.close()
      result = None
  if result is not None:
    # happy days we have a connection
    # I should get here the third time
    pass

The except clause (and it's internal if) I copied from the internals of socket and want to verify that I "test" that path through my copy of the code. (I don't understand how socket can get to that code (setting the target while still raising an exception, but that isn't my concern, only the I verify that I can replicate that code path.) That's why I want the side effect to happen when the mock is called and not when I build the list.

Solution 1

According to the unittest.mock documentation for side_effect:

If you pass in an iterable, it is used to retrieve an iterator which must yield a value on every call. This value can either be an exception instance to be raised, or a value to be returned from the call to the mock (DEFAULT handling is identical to the function case).

Therefore, your socket.create_connection mock will return the function oddConnect for the first two calls, then return the Mock object for the last call. From what I understand, you want to mock create_connection object to actually call those functions as side effects rather than returning them.

I find this behavior rather odd, since you'd expect side_effect, to mean side_effect in every case, not return_value. I suppose the reason this is so lies in the fact that the value of the return_value property must be interpreted as-is. For instance, if your Mock had return_value=[1, 2, 3], would your Mock return [1, 2, 3] for every call, or would it return 1 for the first call?

Solution

Fortunately, there is a solution to this problem. According to the docs, if you pass a single function to side_effect, then that function will be called (not returned) every time the mock is called.

If you pass in a function it will be called with same arguments as the mock and unless the function returns the DEFAULT singleton the call to the mock will then return whatever the function returns. If the function returns DEFAULT then the mock will return its normal value (from the return_value).

Therefore, in order to achieve the desired effect, your side_effect function must do something different every time it is called. You can easily achieve this with a counter and some conditional logic in your function. Note that in order for this to work, your counter must exist outside the scope of the function, so the counter isn't reset when the function exits.

import mock
import socket

# You may wish to encapsulate times_called and oddConnect in a class
times_called = 0
def oddConnect():
  times_called += 1
  # We only do something special the first two times oddConnect is called
  if times_called <= 2:
    result = mock.MagicMock()  # this is where the return value would go
    raise socket.error  # I want it assigned but also this raised  

socket.create_connection = mock(spec=socket.create_connection,
  side_effect=oddConnect)
# what I want: call my function twice, and on the third time return normally
# what I get: two function objects returned and then the normal return

for _ in xrange(3):
  result = None
  try:
    # this is the context in which I want the oddConnect function call
    # to be called (not above when creating the list)
    result = socket.create_connection()
  except socket.error:
    if result is not None:
      # I should get here twice
      result.close()
      result = None
  if result is not None:
    # happy days we have a connection
    # I should get here the third time
    pass

Solution 2

I also encountered the problem of wanting a side effect to occur for only some items in a list of values.

In my case, I wanted to call a method from freezegun the third time my mocked method was called. These answers were really helpful for me; I ended up writing up a fairly general wrapper class, which I thought I'd share here:

class DelayedSideEffect:
    """
    If DelayedSideEffect.side_effect is assigned to a mock.side_effect, allows you to
    delay the first call of callback until after a certain number of iterations.
    """
    def __init__(self, callback, delay_until_call_num: int, return_value=DEFAULT):
        self.times_called = 0
        self.delay_until_call_num = delay_until_call_num
        self.callback = callback
        self.return_value = return_value

    def side_effect(self, *args, **kwargs):
        self.times_called += 1
        if self.times_called >= self.delay_until_call_num:
            self.callback()
        return self.return_value

To then return "my_default_return_value" without calling the lambda function on the first three calls:

with freeze_time(datetime.now()) as freezer:
    se = DelayedSideEffect(callback=lambda: freezer.move_to(the_future), 3)
    my_mock = MagicMock(return_value="my_default_return_value", side_effect=se)