Async-aware Mock/override Environment Variables?

Is there an async-aware way to mock environment variables?

What I want is a context manager that lets me mock the code inside the with block, so that when that code is executing the mock is applied, but it's not applied when other code is executing.

Something like this, but passing:

import pytest
import os
import asyncio
import mock


async def aaa():
    with mock.patch.dict(os.environ, {"FAVORITE_LETTER": "A"}, clear=True):
        for _ in range(50):
            assert os.environ["FAVORITE_LETTER"] == "A"
            await asyncio.sleep(0)


async def bbb():
    with mock.patch.dict(os.environ, {"FAVORITE_LETTER": "B"}, clear=True):
        for _ in range(50):
            assert os.environ["FAVORITE_LETTER"] == "B"
            await asyncio.sleep(0)


@pytest.mark.asyncio
async def test_mock_env():
    await asyncio.gather(aaa(), bbb())

This currently fails with:

    async def aaa():
        with mock.patch.dict(os.environ, {"FAVORITE_LETTER": "A"}, clear=True):
            for _ in range(50):
>               assert os.environ["FAVORITE_LETTER"] == "A"
E               AssertionError: assert 'B' == 'A'
E                 - A
E                 + B

test_mock_env.py:9: AssertionError

=======

Edit:

The accepted answer addresses the question as I asked it, but doesn't solve my problem because I simplified too much when I made the example. My real problem wouldn't make a good example to paste here, but this is closer to my real situation:

from my_library import some_component_i_control
from someone_elses_library import not_my_code

async def wrap_not_my_code():
    with patch.dict(os.environ, {"ENV_VARS_NEEDED_FOR_COMPONENT_I_DONT_CONTROL": "A"}, clear=True):
        await not_my_code()  # this runs forever but does not block the event loop


@pytest.mark.asyncio
async def test_mock_env():
    task_for_not_my_code = asyncio.create_task(wrap_not_my_code())
    task_for_my_component = asyncio.create_task(some_component_i_control())
    await asyncio.wait([task_for_my_component, task_for_not_my_code], return_when=asyncio.FIRST_COMPLETED) 
    # ...
    # Make some assertions about the results

I'm trying to test the interaction between a library I control, and one I don't, with the component I don't control seeing a modified version of the environment.

In fact it's even worse: I need multiple instances of the task for the component I don't control, all seeing different environments, and all running concurrently with my own code. Is that possible?

You will need something to allow exclusive access to state, in this case an asyncio.Lock object would work. This requires you to modify the original code to the following.

import pytest
import os
import asyncio
from unittest.mock import patch


async def aaa(lock):
    for _ in range(50):
        async with lock:
            with patch.dict(os.environ, {"FAVORITE_LETTER": "A"}, clear=True):
                assert os.environ["FAVORITE_LETTER"] == "A"
        await asyncio.sleep(0)


async def bbb(lock):
    for _ in range(50):
        async with lock:
            with patch.dict(os.environ, {"FAVORITE_LETTER": "B"}, clear=True):
                assert os.environ["FAVORITE_LETTER"] == "B"
        await asyncio.sleep(0)


@pytest.mark.asyncio
async def test_mock_env():
    lock = asyncio.Lock()
    await asyncio.gather(aaa(lock), bbb(lock))