如何使用unittest.mock.patch模拟os.environ

Question

我有一些代码，比如：

def get_room_by_computer_number(metrics, computer: str) -> typing.Optional[str]:
logger.info(f"Searching room by device: {computer}")
if not os.environ.get('STAGE') != 'phi':
return None
# more code here

我正试着这样测试：

@patch("requests.post")
def test_if_room_exist(self, mock_requests):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"data": {
"platform_devices_management_device_assignments": [
{"asset_id": "PL00", "room": "room 1", "serial_number": "255CXP3"}
]
}
}
mock_requests.post.return_value = mock_response
self.assertEqual(get_room_by_computer_number(computer="PL00"), "room 1")

如果未在os.environ中设置特定值，则代码将为return None。我想测试在设置值的情况下会发生什么。

如何模拟os.environ中的数据以通过测试？

Answer 1

>>> from unittest.mock import patch
>>> with patch("os.environ",{"poop":"smells"}):
...    import os
...    print(os.environ)
...
{'poop': 'smells'}
>>> @patch('os.environ',{'foo':'bar'})
... def a_func(x,*a,**kw):
...     import os
...     print(os.environ)
...
>>> a_func(55)
{'foo': 'bar'}
>>>

对我来说似乎很简单。。。

如果你只是想覆盖一些小的

>>> import os
>>> @patch('os.environ',{**os.environ,**{'foo':'bar'}})
... def a_func(x,*a,**kw):
...     print(os.environ)
...
>>> a_func(55)
{'SHELL': '/bin/bash', 'NVM_INC': '/home/s