Question

我正在使用请求创建一个requests.Session并将其设置在一个函数中：

def create_web_session(cookie=None):
    s = requests.Session()
    if cookie is not None:
        s.cookies.set("my_cookie_name", cookie)
    return s

会话可以用作上下文管理器。我可以使用在with语句中返回会话（上方）的函数吗？

with create_web_session("my_cookie_value") as s:
     s.get(...)

或者我是否必须更改函数，该函数需要一个会话并在上下文管理器中进行设置：

with requests.Session() as s:
    setup_web_session(s, "my_cookie_value")
    s.get(...)

运行前者似乎可行，但是我的单元测试失败，这就是为什么我问这个问题。我正在修补create_web_session，并在测试s.get(...)时断言我的模拟被.get()调用了，但是似乎是session_mock.__enter__()在调用它。这是预期的行为吗？

这里是一个例子：

# Function
def my_function():
    s = create_web_session()
    s.get("https://google.com")
    s.close()

# Test
@patch("foo.bar.create_web_session")
def test_my_function(self, mock_create_web_session):
    my_function()
    mock_create_web_session.assert_called_once()
    mock_create_web_session.return_value.get.assert_called_once()
    mock_create_web_session.return_value.close.assert_called_once()

一旦我更改了功能以使用上下文管理器：

def my_function():
    with create_web_session() as s:
        s.get("https://google.com")

测试失败并显示： Expected 'get' to have been called once. Called 0 times.

Answer 1

您的create_web_session很好。测试中的问题是，尽管requests.Session.__enter__仅返回了相同的会话，但是模拟中的所有方法都返回了一个新的模拟对象。我们可以告诉模拟行为按照我们想要的方式进行操作，并获得如下工作示例：

def test_my_function(self, mock_create_web_session):
    session = mock_create_web_session.return_value
    session.__enter__.return_value = session

    my_function()

    mock_create_web_session.assert_called_once()
    session.get.assert_called_once()
    session.__exit__.assert_called_once()

请注意，我断言__exit__是被调用的，而不是close，因为该模拟对close或实际会话一无所知。

我可以在使用之前定义和设置上下文管理器吗？

1 个答案: