How to use mock patch os.environ python

Question

am tryin to write unittest for some functions to test if PC number is Equal to the room number which correct in the code below, but if its not phi os.environ return None and which i get, my question is how do i mock the os.environ 'phi' to pass the test?

@metric_scope
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

        ################.....

@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")
   

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'}
>>>

seems pretty straightforward to me...