Mocking a function call inside a context manager

Viewed 24

I'm writing a unit test for a function that calls a method inside a context manager like so:

a.py:

def foo():
   x = None
   with tarfile.open("file_a.tar", "r") as tar:
      x = tar.getmembers()
   return x

and would like to mock getmembers such that I can set an arbitrary return value for x

I am currently mocking the context manager based off this previous question, to no avail:

from unittest.mock import patch, mock_open, call, MagicMock, Mock

import a

@patch("a.tarfile.open")
def test_foo(mockA):
   expected_res = ["asdf"]
   getm = MagicMock()
   getm.getmembers.return_value = ["asdf"]
   mockA.return_value = Mock(tar = getm, __enter__ = mockA, __exit__ = Mock())

   res = a.foo()

   assert res == expected_res

I'm not sure what I'm doing wrong here, so any help would be greatly appreciated!

0 Answers
Related