Skip to content

Testing Utilities API Reference

API documentation for testing utilities and fixtures.

aio_azure_clients_toolbox.testing_utils.fixtures

async_cosmos()

This is unfortunately complex because the Async Cosmos Client uses an Async context manager which returns a client with sync methods.

Ultimately, we want to make it easy for a test function to mock out what the container client returns, but to do that, we have to mock:

  • CosmosClient (with a context manager)
  • Database Client
  • Container Client

The last one is the one that we really want to make available to callers.

If you get confused (and it's unfortunately easy to be confused here), you should litter print("Container Client", id(container_mock)) all over so you can pin down which mock thing you are looking at.

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture()
def async_cosmos():
    """
    This is unfortunately complex because the Async Cosmos Client uses
    an Async context manager which returns a client with _sync_ methods.

    Ultimately, we want to make it easy for a test function to mock out
    what the container client returns, but to do that, we have to mock:

    - CosmosClient (with a context manager)
    - Database Client
    - Container Client

    The last one is the one that we really want to make available to callers.

    If you get confused (and it's unfortunately easy to be confused here),
    you should litter `print("Container Client", id(container_mock))` all over
    so you can *pin down* _which_ mock thing you are looking at.
    """

    # the parent client
    cosmos_client_mock = mock.MagicMock(az_cosmos_client.CosmosClient)
    cosmos_client_mock.return_value = cosmos_client_mock

    # the database client (in the middle)
    cosmos_db_mock = mock.MagicMock(az_cosmos_client.DatabaseProxy)
    cosmos_client_mock.get_database_client.return_value = cosmos_db_mock

    # the important one where our queries will happen
    container_mock = mock.AsyncMock(az_cosmos_client.ContainerProxy)
    # the thing that will *give us back* the above
    cosmos_db_mock.get_container_client.return_value = container_mock

    return cosmos_client_mock, container_mock

cosmos_client_mock(monkeypatch, async_cosmos)

Mock out Azure Async Cosmos client

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture(autouse=True)
def cosmos_client_mock(monkeypatch, async_cosmos):
    """Mock out Azure Async Cosmos client"""
    cosmos_client_mock, _ = async_cosmos
    monkeypatch.setattr(az_cosmos_client, "CosmosClient", cosmos_client_mock)
    monkeypatch.setattr(clients.cosmos, "CosmosClient", cosmos_client_mock)
    monkeypatch.setattr(clients.cosmos.ManagedCosmos, "ready", mock.AsyncMock())

    return cosmos_client_mock

cosmos_deletable(async_cosmos)

Patch delete_item on container_mock.

Note: this is a lot simpler than returning what delete_item returns.

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture()
def cosmos_deletable(async_cosmos):
    """Patch `delete_item` on container_mock.

    Note: this is a lot simpler than returning what delete_item returns.
    """
    _, container_mock = async_cosmos

    def set_return(return_value, side_effect=None):
        container_mock.delete_item.return_value = return_value
        container_mock.delete_item.side_effect = side_effect

    yield container_mock, set_return

cosmos_insertable(async_cosmos)

Patch create_item on container_mock.

Note: this is a lot simpler than returning what create_item returns.

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture()
def cosmos_insertable(async_cosmos):
    """Patch `create_item` on container_mock.

    Note: this is a lot simpler than returning what create_item returns.
    """
    _, container_mock = async_cosmos

    def set_return(return_value, side_effect=None):
        container_mock.create_item.return_value = return_value
        container_mock.create_item.side_effect = side_effect

    yield container_mock, set_return

cosmos_patchable(async_cosmos)

Patch patch_item on container_mock.

Note: this is a lot simpler than returning what read_item returns.

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture()
def cosmos_patchable(async_cosmos):
    """Patch `patch_item` on container_mock.

    Note: this is a lot simpler than returning what read_item returns.
    """
    _, container_mock = async_cosmos

    def set_return(return_value, side_effect=None):
        container_mock.patch_item.return_value = return_value
        container_mock.patch_item.side_effect = side_effect

    yield container_mock, set_return

cosmos_queryable(async_cosmos)

Make it possible to query a specific result back from the Cosmos mock.

Example

import itertools

_, set_return = cosmos_upsertable

return a thing ONE TIME (remember: this is an async iterator)

set_return("hello")

or repeatedly return the same thing

set_return(None, side_effect=itertools.repeat("hello"))

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture()
def cosmos_queryable(async_cosmos):
    """Make it possible to query a specific result back from the Cosmos mock.

    Example:
        import itertools

        _, set_return = cosmos_upsertable
        # return a thing ONE TIME (remember: this is an async iterator)
        set_return("hello")

        # or repeatedly return the same thing
        set_return(None, side_effect=itertools.repeat("hello"))

    """
    _, container_mock = async_cosmos

    def set_return(return_value, side_effect=None):
        async_iter_thing = AsyncIterImplementation()
        async_iter_thing.return_value = return_value
        container_mock.query_items.return_value = async_iter_thing
        if side_effect is not None:
            async_iter_thing.side_effect = side_effect

    yield container_mock, set_return

cosmos_readable(async_cosmos)

Patch read_item on container_mock.

Note: this is a lot simpler than returning what read_item returns.

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture()
def cosmos_readable(async_cosmos):
    """Patch `read_item` on container_mock.

    Note: this is a lot simpler than returning what read_item returns.
    """
    _, container_mock = async_cosmos

    def set_return(return_value, side_effect=None):
        container_mock.read_item.return_value = return_value
        container_mock.read_item.side_effect = side_effect

    yield container_mock, set_return

cosmos_updatable(async_cosmos)

Patch replace_item on container_mock.

Note: this is a lot simpler than returning what replace_item returns.

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture()
def cosmos_updatable(async_cosmos):
    """Patch `replace_item` on container_mock.

    Note: this is a lot simpler than returning what replace_item returns.
    """
    _, container_mock = async_cosmos

    def set_return(return_value, side_effect=None):
        container_mock.replace_item.return_value = return_value
        container_mock.replace_item.side_effect = side_effect

    yield container_mock, set_return

cosmos_upsertable(async_cosmos)

Patch patch_item on container_mock.

Note: this is a lot simpler than returning what read_item returns.

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture()
def cosmos_upsertable(async_cosmos):
    """Patch `patch_item` on container_mock.

    Note: this is a lot simpler than returning what read_item returns.
    """
    _, container_mock = async_cosmos

    def set_return(return_value, side_effect=None):
        container_mock.upsert_item.return_value = return_value
        container_mock.upsert_item.side_effect = side_effect

    yield container_mock, set_return

mock_azureblob(monkeypatch) async

Mock out Azure Blob Service client and its children.

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture(autouse=True)
async def mock_azureblob(monkeypatch):  # type: ignore
    """Mock out Azure Blob Service client and its children."""
    bsc = mock.MagicMock(BlobServiceClient)
    container_client = mock.MagicMock(ContainerClient)
    mockblobc = mock.MagicMock(BlobClient)

    # Parent: gets container client, delivers a blob client
    # Used for get-lease: overwrite if needed
    bsc.account_name = "our-company-blobs"
    bsc.get_container_client = mock.Mock(return_value=container_client)
    bsc.get_blob_client = mock.PropertyMock(return_value=mockblobc)
    container_client.get_blob_client = mock.PropertyMock(return_value=mockblobc)

    # Setattr the parent for this codebase
    monkeypatch.setattr(
        clients.azure_blobs, "BlobServiceClient", make_async_ctx_mgr(bsc)
    )

    def set_download_return(return_value, side_effect=None):
        """
        Set the return value for the download_blob method.
        """
        # For `async for` with `download_blob`
        async_downloader = AsyncIterImplementation()
        async_downloader.return_value = return_value

        def make_chunks(value):
            return mock.Mock(
                **{
                    "readall": mock.AsyncMock(return_value=value),
                    "chunks.return_value": async_downloader,
                }
            )

        chunks_thing = make_chunks(return_value)
        mockblobc.download_blob.return_value = chunks_thing

        # hack download_blob to return different things on each call
        if side_effect is not None and isinstance(side_effect, list):
            all_chunks = [make_chunks(v) for v in side_effect]
            mockblobc.download_blob.side_effect = all_chunks
        elif side_effect is not None:
            async_downloader.side_effect = side_effect

    def set_list_blobs_return(list_or_side_effect):
        # For `async for` with `list_blobs`
        async_lister = AsyncIterImplementation()
        async_lister.side_effect = list_or_side_effect
        container_client.list_blobs.return_value = async_lister

    class SetReturns:
        def __init__(self):
            self.download_blob_returns = set_download_return
            self.list_blobs_returns = set_list_blobs_return

    return container_client, mockblobc, SetReturns()

mockegrid(monkeypatch)

Mock out Azure Eventgrid client

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture(autouse=True)
def mockegrid(monkeypatch):  # type: ignore
    """Mock out Azure Eventgrid client"""
    mockeg = mock.MagicMock(spec=EventGridPublisherClient)
    mockeg.return_value = mockeg

    async_mockeg = mock.MagicMock(spec=AsyncEventGridPublisherClient)
    async_mockeg.return_value = make_async_ctx_mgr(async_mockeg)

    monkeypatch.setattr(clients.eventgrid, "EventGridPublisherClient", mockeg)
    monkeypatch.setattr(clients.eventgrid, "AsyncEventGridPublisherClient", async_mockeg)
    return mockeg, async_mockeg

mockehub(monkeypatch)

Mock out Azure Blob Service client

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture(autouse=True)
def mockehub(monkeypatch):  # type: ignore
    """Mock out Azure Blob Service client"""
    mockev = mock.MagicMock(spec=EventHubProducerClient)
    mockev.create_batch = mock.AsyncMock(return_value=mockev)
    mockev.send_batch = mock.AsyncMock()
    mockev.add = mock.Mock()

    def get_client(*args, **kwargs):
        return mockev

    monkeypatch.setattr(clients.eventhub.Eventhub, "get_client", get_client)
    return mockev

mockservicebus(monkeysession)

Mock out Azure Blob Service client

Source code in aio_azure_clients_toolbox/testing_utils/fixtures.py
@pytest.fixture(autouse=True)
def mockservicebus(monkeysession):  # type: ignore
    """Mock out Azure Blob Service client"""
    mocksb = mock.MagicMock(ServiceBusClient)

    # Receiver client
    receiver_client = mock.AsyncMock()
    ctx_receiver = make_async_ctx_mgr(receiver_client)
    mocksb.get_queue_receiver = mock.PropertyMock(return_value=ctx_receiver)
    # Sender client reuse mocksb
    sender_client = mock.AsyncMock()
    ctx_sender = make_async_ctx_mgr(sender_client)
    mocksb.get_queue_sender = mock.PropertyMock(return_value=ctx_sender)

    def get_client(*args, **kwargs):
        # Wrap it up so it can be used as a context manager
        return make_async_ctx_mgr(mocksb)

    monkeysession.setattr(
        clients.service_bus,
        "ServiceBusClient",
        get_client,
    )
    mocksb._receiver = receiver_client
    mocksb._sender = sender_client
    sender_client.schedule_messages = mock.AsyncMock()

    return mocksb