How do I write a proper test for this mocked context manager and ZipFile?

Question

How am I supposed to write this test? I've tried the various options listed but each returns the same failed test:

import zipfile
from mock import Mock, patch    

def unzip_file(fp):
    with zipfile.ZipFile(fp, 'r') as z:
        z.extractall('dir')
    
    
@patch('zipfile.ZipFile')
def test_unzip_file(m_zipfile):

    # I have tried the following...
    
    # m_zipfile.__enter__.extractall = Mock()
    # m_zipfile.extractall = Mock()
    # m_zipfile.return_value.__enter__.return_value = Mock()
    m_zipfile.return_value.__enter__.return_value.extractall = Mock()

    unzip_file('test')

    m_zipfile.assert_called_with('test', 'r')  # this test passes
    m_zipfile.extractall.assert_called_with('dir')  # this test fails

I tried to use this answer as a guide but I'm still lost as to how to properly do this. The actual function in our code is more complex with additional parameters but I am trying to start at the base first.

The failure...

E           AssertionError: expected call not found.
E           Expected: extractall('dir')
E           Actual: not called.

Answer 1

You're pretty close, you just need a reference to the mock within the context manager to assert against

@patch('zipfile.ZipFile')
def test_unzip_file(m_zipfile):
    m_extractall = Mock()
    m_zipfile.return_value.__enter__.return_value.extractall = m_extractall

    unzip_file('test')

    m_zipfile.assert_called_with('test', 'r')  # this test passes
    m_extractall.assert_called_with('dir') # and this one does, too