<div dir="ltr"><div><div>Hi<br><br></div>There was a discussion about a similar question back in 2010 [1], but I am not aware of any solution.<br><br></div><div>Suppose you have a.py and test.py<br><br>#a.py<br>import contextlib<br><br>@contextlib.contextmanager<br>def foo(a):<br>    yield &quot;123&quot;<br><br>def bar():<br>    with foo(1) as f:<br>        return f<br><br></div><div>#test.py<br>import unittest<br>from mock import patch, Mock, MagicMock, mock_open<br><br></div><div>import a<br></div><div><br>class TestFoo(unittest.TestCase):<br>    def test_foo(self):<br>        with patch(&quot;a.foo&quot;, new=MagicMock) as f:<br>            f.return_value = &quot;456&quot;<br>            self.assertEqual(a.bar(), &quot;456&quot;)<br><br>if __name__ == &quot;__main__&quot;:<br>    unittest.main()<br></div><div><br></div><div>This won&#39;t work as foo takes an argument.  I am getting the following exception:<br>======================================================================<br>ERROR: test_foo (__main__.TestFoo)<br>----------------------------------------------------------------------<br>Traceback (most recent call last):<br>  File &quot;test1.py&quot;, line 9, in test_foo<br>    self.assertEqual(a.bar(), &quot;456&quot;)<br>  File &quot;/private/tmp/a.py&quot;, line 8, in bar<br>    with foo(1) as f:<br>AttributeError: __exit__<br><br>I tried to first patch the context manager decorator, and then patch a.foo (suggested in [2]), but I am still getting the same exact error. If I don&#39;t have the argument, the code above (using new=MagicMock) will work.<br><br>Has anyone came up with a solution for this problem yet? I can&#39;t imagine I am the only one having this issue.<br><br></div><div>Thank you.<br></div><div><br>[1]: <a href="http://lists.idyll.org/pipermail/testing-in-python/2010-December/003703.html">http://lists.idyll.org/pipermail/testing-in-python/2010-December/003703.html</a><br></div>[2]: <a href="http://jamescooke.info/things-to-remember-about-decorators.html">http://jamescooke.info/things-to-remember-about-decorators.html</a><br></div>