Testing response status for aiohttp ClientSession

test_response_status.py

import random
from aiohttp import ClientSession
from asynctest import CoroutineMock, MagicMock, patch


async def get_random_photo_url():
    while True:
        async with ClientSession() as session:
            async with session.get('random.photos') as resp:
                json = await resp.json()
        photos = json['photos']
        if not photos or resp.status != 200:
            continue
        return random.choice(photos)['img_src']

# TESTS


class AsyncContextManagerMock(MagicMock):
    async def __aenter__(self):
        return self.aenter()

    async def __aexit__(self, *args):
        pass


class MockedResponse(MagicMock):
    def __init__(self, status, json=None, *args):
        super().__init__(*args)
        self.status = status
        self.json = CoroutineMock(return_value=json)


@patch('aiohttp.ClientSession.get', new_callable=AsyncContextManagerMock)
async def test_status(mock_get):
    mock_get.return_value.aenter.side_effect = [
        MockedResponse(500, {'photos': [{'img_src': 'a.jpg'}]}),
        MockedResponse(200, {'photos': [{'img_src': 'b.jpg'}]}),
    ]

    image_url = await get_random_photo_url()

    assert mock_get.return_value.aenter.call_count == 2
    assert image_url == 'b.jpg'