Mocking functions from object created by context manager

Question

I have the following module that I am trying to write unit tests for.

import myModuleWithCtxMgr

def myFunc(arg1):
    with myModuleWithCtxMgr.ctxMgr() as ctxMgr:
        result = ctxMgr.someFunc()

    if result:
        return True, result
    return False, None

The unit tests I'm working on looks like this.

import mock
import unittest
import myModule as myModule

class MyUnitTests(unittest.TestCase):

    @mock.patch("myModuleWithCtxMgr.ctxMgr")
    def testMyFunc(self, mockFunc):
        mockReturn = mock.MagicMock()
        mockReturn.someFunc = mock.Mock(return_value="val")
        mockFunc.return_value = mockReturn
        result = myModule.myFunc("arg")

        self.assertEqual(result, (True, "val"))

The test is failing because result[0] = magicMock() and not the return value (I thought) I configured.

I've tried a few different variations of the test but I can't seem to be able to mock the return value of ctxMgr.someFunc(). Does anyone know how I might accomplish this?

Thanks!

Answer 1

The error says:

First differing element 1:
<MagicMock name='ctxMgr().__enter__().someFunc()' id='139943278730000'>
'val'

- (True, <MagicMock name='ctxMgr().__enter__().someFunc()' id='139943278730000'>)
+ (True, 'val')

The error contains the mock name which exactly shows you what needs to be mocked. Note that __enter__ corresponds to the Context Manager protocol.

This works for me:

class MyUnitTests(unittest.TestCase):

    @mock.patch("myModuleWithCtxMgr.ctxMgr")
    def testMyFunc(self, mockCtxMgr):
        mockCtxMgr().__enter__().someFunc.return_value = "val"
        result = myModule.myFunc("arg")

        self.assertEqual(result, (True, "val"))

Note how each of these is a separate MagicMock instance which you can configure:

• mockCtxMgr
• mockCtxMgr()
• mockCtxMgr().__enter__
• mockCtxMgr().__enter__()
• mockCtxMgr().__enter__().someFunc

MagicMocks are created lazily but have identity, so you can configure them this way and it Just Works.