Mock requests.Session.get in a nested context manager

Question

I want to mock a function that uses a requests.Session object, which was created in a context manager, to perform a get request, again using a context manager. So there are two context managers involved.

Here's an example code:

# main.py

from requests import Session

def fun(session):
    with session.get("https://httpstat.us/200") as response:
        print(response)

        
def run():
    with Session() as session:
        fun(session)

        
if __name__ == "__main__":
    run()

When executed, the status code of 200 is printed as expected.

Now I want to mock the get request, to return a different value. Somehow I can't navigate the mess of context managers and return values.

Here's my best effort:

# maintest.py
from unittest.mock import patch
from main import run

def test_main():
    
    with patch("main.Session") as mocked_session:
        mocked_session.__enter__.return_value.get.__enter__.return_value = "MOCKED"
        run()

Running the test with pytest, I was hoping to get "MOCKED" printed, but instead I get

<MagicMock name='Session().__enter__().get().__enter__()' id='140303100714240'>

I've tried any conceivable combination of __enter__, .get and .return_value

Answer 1

You got it almost right, the correct way to mock is:

with patch("main.Session") as mocked_session:
    mocked_session.return_value.__enter__.return_value.get.return_value.__enter__.return_value = "MOCKED"

I will break this down for clarification.

Your session_mock refers to the Session class. To get the Session instance, you use return_value (as you do for every call).
The with statement internally calls __enter__() on that session to get the session instance, so you have to add __enter__ for the method and return_value for the call.
Now you call get() on that object, which by the same logic needs an added get.return_value.
And finally, this is used in another with statement to create response, which again needs an added __enter__.return_value.

The result gets you the response object for your mocked Session, and you can set another value.

If breaking that down in code, it would look something like:

with patch("main.Session") as mocked_session:
    mocked_session_instance = mocked_session.return_value
    mocked_session_object = mocked_session_instance.__enter__.return_value
    mocked_get = mocked_session_object.get.return_value
    mocked_response = mocked_get.__enter__
    mocked_response.return_value = "MOCKED"