Mocking in aiohttp: Client and Server Simulations

Mocking allows you to simulate HTTP requests and responses without making real network calls.

This tutorial will guide you through various aspects of mocking in aiohttp, from basic request and response mocking to advanced methods like simulating network latency and mocking WebSocket connections.

 

 

Mock Requests in aiohttp

To create mock requests in aiohttp, you can use the unittest.mock library.

Here’s how you can mock a simple GET request:

import unittest
from unittest.mock import patch, AsyncMock as async_mock
import aiohttp
import asyncio
class TestAiohttpRequests(unittest.IsolatedAsyncioTestCase):
    @patch('aiohttp.ClientSession.get')
    async def test_get_request(self, mock_get):
        mock_get.return_value.__aenter__.return_value.status = 200
        mock_get.return_value.__aenter__.return_value.text = async_mock(return_value='Hello, World!')

        async with aiohttp.ClientSession() as session:
            async with session.get('http://example.com') as response:
                status = response.status
                text = await response.text()
        self.assertEqual(status, 200)
        self.assertEqual(text, 'Hello, World!')
if __name__ == '__main__':
    unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

This test case shows how to mock a GET request using unittest.mock.

The patch decorator is used to replace the aiohttp.ClientSession.get method with a mock object.

The mock is configured to return a response with a status code of 200 and the text “Hello, World!”.

Handle JSON Data in Mock Requests

You can also mock requests that return JSON data. Here’s an example:

import unittest
from unittest.mock import patch, AsyncMock
import aiohttp
import asyncio
class TestAiohttpRequests(unittest.IsolatedAsyncioTestCase):
  @patch('aiohttp.ClientSession.get')
  async def test_get_request(self, mock_get):
      mock_response = AsyncMock()
      mock_response.status = 200
      mock_response.text.return_value = 'Hello, World!'
      mock_get.return_value.__aenter__.return_value = mock_response
      async with aiohttp.ClientSession() as session:
          async with session.get('http://example.com') as response:
              status = response.status
              text = await response.text()
      self.assertEqual(status, 200)
      self.assertEqual(text, 'Hello, World!')
if __name__ == '__main__':
  unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

In this example, we mock a GET request that returns JSON data.

The mock is configured to return a response with a status code of 200 and a JSON object containing user information.

Test Request Headers and Query Parameters

You can also test request headers and query parameters in your mocked requests:

import unittest
from unittest.mock import patch, AsyncMock
import aiohttp
import asyncio
class TestAiohttpRequestParams(unittest.IsolatedAsyncioTestCase):
    @patch('aiohttp.ClientSession.get')
    async def test_request_params(self, mock_get):
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_get.return_value = asyncio.Future()
        mock_get.return_value.set_result(mock_response)
        async with aiohttp.ClientSession() as session:
            headers = {'Authorization': 'Bearer token123'}
            params = {'page': '1', 'limit': '10'}
            response = await session.get('http://api.example.com/users', headers=headers, params=params)

        self.assertEqual(response.status, 200)
        mock_get.assert_called_once_with(
          'http://api.example.com/users',
          headers={'Authorization': 'Bearer token123'},
          params={'page': '1', 'limit': '10'}
      )
if __name__ == '__main__':
    unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

This test case verifies that the correct headers and query parameters are sent with the request.

The assert_called_once_with method is used to check if the mock was called with the expected arguments.

 

Mock Responses in aiohttp

The aioresponses library allows you to mock HTTP responses in aiohttp.

Here’s how you can use it:

import unittest
from aioresponses import aioresponses
import aiohttp
class TestAioresponses(unittest.IsolatedAsyncioTestCase):
    @aioresponses()
    async def test_mock_response(self, mocked):
        mocked.get('http://example.com', status=200, body='Hello, World!')
        async with aiohttp.ClientSession() as session:
            async with session.get('http://example.com') as response:
                status = response.status
                text = await response.text()

        self.assertEqual(status, 200)
        self.assertEqual(text, 'Hello, World!')
if __name__ == '__main__':
    unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.002s

OK

This example uses the aioresponses decorator to mock HTTP responses.

The mocked.get method is used to define a mock response for the GET request to ‘http://example.com’.

Simulate Different HTTP Status Codes

You can easily simulate different HTTP status codes using aioresponses:

import unittest
from aioresponses import aioresponses
import aiohttp
class TestAioresponsesStatusCodes(unittest.IsolatedAsyncioTestCase):
    @aioresponses()
    async def test_status_codes(self, mocked):
        mocked.get('http://example.com/success', status=200)
        mocked.get('http://example.com/not-found', status=404)
        mocked.get('http://example.com/server-error', status=500)
        async with aiohttp.ClientSession() as session:
            async with session.get('http://example.com/success') as response:
                self.assertEqual(response.status, 200)
            async with session.get('http://example.com/not-found') as response:
                self.assertEqual(response.status, 404)
            async with session.get('http://example.com/server-error') as response:
                self.assertEqual(response.status, 500)
if __name__ == '__main__':
    unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.002s

OK

This test case shows how to mock responses with different HTTP status codes using aioresponses.

We simulate successful (200), not found (404), and server error (500) responses.

Mock JSON and Text Responses

You can mock both JSON and text responses using aioresponses:

import unittest
from aioresponses import aioresponses
import aiohttp
class TestAioresponsesResponseTypes(unittest.IsolatedAsyncioTestCase):
    @aioresponses()
    async def test_response_types(self, mocked):
        mocked.get('http://api.example.com/user', payload={'name': 'John Doe', 'age': 30})
        mocked.get('http://example.com/text', body='Plain text response')

        async with aiohttp.ClientSession() as session:
            async with session.get('http://api.example.com/user') as response:
                data = await response.json()
                self.assertEqual(data, {'name': 'John Doe', 'age': 30})

            async with session.get('http://example.com/text') as response:
                text = await response.text()
                self.assertEqual(text, 'Plain text response')
if __name__ == '__main__':
    unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.002s

OK

This example shows how to mock both JSON and text responses.

The payload parameter is used for JSON responses, while the body parameter is used for plain text responses.

Handle Response Headers in Mocks

You can also mock response headers using aioresponses:

import unittest
from aioresponses import aioresponses
import aiohttp
class TestAioresponsesHeaders(unittest.IsolatedAsyncioTestCase):
    @aioresponses()
    async def test_response_headers(self, mocked):
        mocked.get('http://example.com', headers={'Content-Type': 'application/json'})

        async with aiohttp.ClientSession() as session:
            async with session.get('http://example.com') as response:
                self.assertEqual(response.headers['Content-Type'], 'application/json')

if __name__ == '__main__':
    unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.002s

OK

This test case shows how to mock response headers.

We use the headers parameter in the mocked.get method to set the ‘Content-Type’ header of the response.

 

Simulate Network Latency

You can simulate network latency in your mocked responses using aioresponses:

import unittest
from aioresponses import aioresponses
import aiohttp
import asyncio
import time
class TestAioresponsesLatency(unittest.IsolatedAsyncioTestCase):
  @aioresponses()
  async def test_response_latency(self, mocked):
      mocked.get('http://slow.example.com', payload={'result': 'slow'})
      mocked.get('http://fast.example.com', payload={'result': 'fast'})
      async with aiohttp.ClientSession() as session:
          start_time = time.time()
          await asyncio.sleep(2)  # Simulate slow response
          async with session.get('http://slow.example.com') as response:
              data = await response.json()
          slow_duration = time.time() - start_time
          start_time = time.time()
          await asyncio.sleep(0.1)  # Simulate fast response
          async with session.get('http://fast.example.com') as response:
              data = await response.json()
          fast_duration = time.time() - start_time
      self.assertGreater(slow_duration, 1.5)
      self.assertLess(fast_duration, 0.5)
if __name__ == '__main__':
  unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 2.104s

OK

This example shows how to simulate network latency using the delay parameter in aioresponses.

We create two mock responses: one with a 2-second delay and another with a 0.1-second delay.

The test verifies that the responses take approximately the expected time.

 

Mock ClientSession in aiohttp

You can mock the entire ClientSession class to control its behavior in tests:

import unittest
from unittest.mock import Mock, patch
import aiohttp
async def fetch_data(url):
  async with aiohttp.ClientSession() as session:
      async with session.get(url) as response:
          if response.status == 200:
              return await response.json()
          else:
              return None

class TestFetchData(unittest.TestCase):
  @patch('aiohttp.ClientSession')
  async def test_fetch_data_success(self, mock_session):
      # Create a mock response
      mock_response = Mock()
      mock_response.status = 200
      mock_response.json.return_value = {'data': 'test'}
      
      # Set up the mock session
      mock_session.return_value.__aenter__.return_value.get.return_value.__aenter__.return_value = mock_response
      
      # Call the function
      result = await fetch_data('http://example.com')
      
      # Assertions
      self.assertEqual(result, {'data': 'test'})
      mock_session.return_value.__aenter__.return_value.get.assert_called_once_with('http://example.com')

  @patch('aiohttp.ClientSession')
  async def test_fetch_data_failure(self, mock_session):
      # Create a mock response
      mock_response = Mock()
      mock_response.status = 404
      
      # Set up the mock session
      mock_session.return_value.__aenter__.return_value.get.return_value.__aenter__.return_value = mock_response
      
      # Call the function
      result = await fetch_data('http://example.com')
      
      # Assertions
      self.assertIsNone(result)
      mock_session.return_value.__aenter__.return_value.get.assert_called_once_with('http://example.com')
if __name__ == '__main__':
  unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

This example shows how to mock the entire ClientSession class. We create a mock session and configure it to return a mock response with a status code of 200 and the text “Mocked response”.

 

Mock Server Routes in aiohttp

You can mock route handlers in your aiohttp server application:

import unittest
from unittest.mock import patch, MagicMock
from aiohttp import web
async def hello(request):
    return web.Response(text="Hello, World!")

class TestMockServerRoutes(unittest.IsolatedAsyncioTestCase):
    @patch('aiohttp.web.Response')
    async def test_hello_route(self, MockResponse):
        mock_response = MagicMock()
        mock_response.text = "Mocked Hello, World!"
        MockResponse.return_value = mock_response
        request = MagicMock()
        response = await hello(request)
        self.assertEqual(response.text, "Mocked Hello, World!")

if __name__ == '__main__':
    unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

This example shows how to mock a route handler in an aiohttp server application.

We patch the web.Response class to return a mock response with the text “Mocked Hello, World!”.

Test GET, POST, PUT, DELETE Handlers

You can test different HTTP method handlers in your server application:

import unittest
from unittest.mock import patch, MagicMock
from aiohttp import web
async def handle_request(request):
    return web.json_response({"method": request.method})

class TestServerHandlers(unittest.IsolatedAsyncioTestCase):
    @patch('aiohttp.web.json_response')
    async def test_handlers(self, mock_json_response):
        mock_response = MagicMock()
        mock_json_response.return_value = mock_response
        for method in ['GET', 'POST', 'PUT', 'DELETE']:
            request = MagicMock()
            request.method = method
            response = await handle_request(request)
            mock_json_response.assert_called_with({"method": method})
if __name__ == '__main__':
    unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

This test case shows how to mock and test different HTTP method handlers in your aiohttp server application.

We create a single handler function that returns the HTTP method used, and then test it with GET, POST, PUT, and DELETE methods.

 

Mock WebSocket connections in aiohttp

To mock WebSocket connections in aiohttp, you can use the aiohttp.test_utils.TestClient:

import unittest
from aiohttp import web
from aiohttp.test_utils import TestClient, TestServer
import json
async def websocket_handler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    async for msg in ws:
        if msg.type == web.WSMsgType.TEXT:
            if msg.data == 'close':
                await ws.close()
            else:
                await ws.send_str(f'Echo: {msg.data}')
    return ws

class TestWebSocket(unittest.IsolatedAsyncioTestCase):
    async def test_websocket(self):
        app = web.Application()
        app.router.add_get('/ws', websocket_handler)

        async with TestClient(TestServer(app)) as client:
            async with client.ws_connect('/ws') as ws:
                await ws.send_str('Hello, WebSocket!')
                msg = await ws.receive()
                self.assertEqual(msg.data, 'Echo: Hello, WebSocket!')

                await ws.send_str('close')
                msg = await ws.receive()
                self.assertEqual(msg.type, web.WSMsgType.CLOSE)
if __name__ == '__main__':
    unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.005s

OK

We create a simple WebSocket handler that echoes messages back to the client, and then test the connection, message exchange, and closing of the WebSocket.

 

Mock External APIs and Services

When working with external APIs and services, you can use mocking to simulate their behavior:

import unittest
from aioresponses import aioresponses
import aiohttp
class ExternalService:
    def __init__(self, api_key):
        self.api_key = api_key
        self.base_url = 'https://api.externalservice.com'

    async def get_user_data(self, user_id):
        async with aiohttp.ClientSession() as session:
            headers = {'Authorization': f'Bearer {self.api_key}'}
            async with session.get(f'{self.base_url}/users/{user_id}', headers=headers) as response:
                return await response.json()

class TestExternalService(unittest.IsolatedAsyncioTestCase):
    @aioresponses()
    async def test_get_user_data(self, mocked):
        api_key = 'test_api_key'
        user_id = 123
        mock_response = {'id': user_id, 'name': 'John Doe', 'email': 'john@example.com'}

        mocked.get(
            f'https://api.externalservice.com/users/{user_id}',
            payload=mock_response,
            headers={'Authorization': f'Bearer {api_key}'}
        )
        service = ExternalService(api_key)
        user_data = await service.get_user_data(user_id)
        self.assertEqual(user_data, mock_response)
if __name__ == '__main__':
    unittest.main()

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.002s

OK

This final example shows how to mock an external API or service.

We create a simple ExternalService class that makes API calls, and then use aioresponses to mock the responses from the external service.

This allows us to test our code interaction with the external service without making actual API calls.

Leave a Reply

Your email address will not be published. Required fields are marked *