Skip to content

Subscriber Testing#

Testability is a crucial part of any application, and FastStream provides you with the tools to test your code easily.

Original Application#

Let's take a look at the original application to test

annotation_kafka.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
from faststream import FastStream
from faststream.kafka import KafkaBroker

broker = KafkaBroker("localhost:9092")
app = FastStream(broker)


@broker.subscriber("test-topic")
async def handle(name: str, user_id: int):
    assert name == "John"
    assert user_id == 1
annotation_rabbit.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
from faststream import FastStream
from faststream.rabbit import RabbitBroker

broker = RabbitBroker("amqp://guest:guest@localhost:5672/")
app = FastStream(broker)


@broker.subscriber("test-queue")
async def handle(name: str, user_id: int):
    assert name == "John"
    assert user_id == 1
annotation_rabbit.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
from faststream import FastStream
from faststream.nats import NatsBroker

broker = NatsBroker("nats://localhost:4222")
app = FastStream(broker)


@broker.subscriber("test-subject")
async def handle(name: str, user_id: int):
    assert name == "John"
    assert user_id == 1

It consumes JSON messages like { "name": "username", "user_id": 1 }

You can test your consume function like a regular one, for sure:

@pytest.mark.asyncio
async def test_handler():
    await handle("John", 1)

But if you want to test your function closer to your real runtime, you should use the special FastStream test client.

In-Memory Testing#

Deploying a whole service with a Message Broker is a bit too much just for testing purposes, especially in your CI environment. Not to mention the possible loss of messages due to network failures when working with real brokers.

For this reason, FastStream has a special TestClient to make your broker work in InMemory mode.

Just use it like a regular async context manager - all published messages will be routed in-memory (without any external dependencies) and consumed by the correct handler.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
import pytest
from pydantic import ValidationError

from faststream.kafka import TestKafkaBroker

from .annotation_kafka import broker, handle


@pytest.mark.asyncio
async def test_handle():
    async with TestKafkaBroker(broker) as br:
        await br.publish({"name": "John", "user_id": 1}, topic="test-topic")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
import pytest
from pydantic import ValidationError

from faststream.rabbit import TestRabbitBroker

from .annotation_rabbit import broker, handle


@pytest.mark.asyncio
async def test_handle():
    async with TestRabbitBroker(broker) as br:
        await br.publish({"name": "John", "user_id": 1}, queue="test-queue")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
import pytest
from pydantic import ValidationError

from faststream.nats import TestNatsBroker

from .annotation_nats import broker, handle


@pytest.mark.asyncio
async def test_handle():
    async with TestNatsBroker(broker) as br:
        await br.publish({"name": "John", "user_id": 1}, subject="test-subject")

Catching Exceptions#

This way you can catch any exceptions that occur inside your handler:

1
2
3
4
5
@pytest.mark.asyncio
async def test_validation_error():
    async with TestKafkaBroker(broker) as br:
        with pytest.raises(ValidationError):
            await br.publish("wrong message", topic="test-topic")

1
2
3
4
5
@pytest.mark.asyncio
async def test_validation_error():
    async with TestRabbitBroker(broker) as br:
        with pytest.raises(ValidationError):
            await br.publish("wrong message", queue="test-queue")

1
2
3
4
5
@pytest.mark.asyncio
async def test_validation_error():
    async with TestNatsBroker(broker) as br:
        with pytest.raises(ValidationError):
            await br.publish("wrong message", subject="test-subject")

Validates Input#

Also, your handler has a mock object to validate your input or call counts.

1
2
3
4
5
6
@pytest.mark.asyncio
async def test_handle():
    async with TestKafkaBroker(broker) as br:
        await br.publish({"name": "John", "user_id": 1}, topic="test-topic")

        handle.mock.assert_called_once_with({"name": "John", "user_id": 1})

1
2
3
4
5
6
@pytest.mark.asyncio
async def test_handle():
    async with TestRabbitBroker(broker) as br:
        await br.publish({"name": "John", "user_id": 1}, queue="test-queue")

        handle.mock.assert_called_once_with({"name": "John", "user_id": 1})

1
2
3
4
5
6
@pytest.mark.asyncio
async def test_handle():
    async with TestNatsBroker(broker) as br:
        await br.publish({"name": "John", "user_id": 1}, subject="test-subject")

        handle.mock.assert_called_once_with({"name": "John", "user_id": 1})

Note

The Handler mock has a not-serialized JSON message body. This way you can validate the incoming message view, not python arguments.

Thus our example checks not mock.assert_called_with(name="John", user_id=1), but mock.assert_called_with({ "name": "John", "user_id": 1 }).

You should be careful with this feature: all mock objects will be cleared when the context manager exits.

1
2
3
4
5
6
7
8
@pytest.mark.asyncio
async def test_handle():
    async with TestKafkaBroker(broker) as br:
        await br.publish({"name": "John", "user_id": 1}, topic="test-topic")

        handle.mock.assert_called_once_with({"name": "John", "user_id": 1})

    assert handle.mock is None

1
2
3
4
5
6
7
8
@pytest.mark.asyncio
async def test_handle():
    async with TestRabbitBroker(broker) as br:
        await br.publish({"name": "John", "user_id": 1}, queue="test-queue")

        handle.mock.assert_called_once_with({"name": "John", "user_id": 1})

    assert handle.mock is None

1
2
3
4
5
6
7
8
@pytest.mark.asyncio
async def test_handle():
    async with TestNatsBroker(broker) as br:
        await br.publish({"name": "John", "user_id": 1}, subject="test-subject")

        handle.mock.assert_called_once_with({"name": "John", "user_id": 1})

    assert handle.mock is None

Real Broker Testing#

If you want to test your application in a real environment, you shouldn't have to rewrite all you tests: just pass with_real optional parameter to your TestClient context manager. This way, TestClient supports all the testing features but uses an unpatched broker to send and consume messages.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import pytest
from pydantic import ValidationError

from faststream.kafka import TestKafkaBroker

from .pydantic_fields_kafka import broker, handle


@pytest.mark.asyncio
async def test_handle():
    async with TestKafkaBroker(broker, with_real=True) as br:
        await br.publish({"name": "John", "user_id": 1}, topic="test-topic")
        await handle.wait_call(timeout=3)
        handle.mock.assert_called_once_with({"name": "John", "user_id": 1})

    assert handle.mock is None

@pytest.mark.asyncio
async def test_validation_error():
    async with TestKafkaBroker(broker, with_real=True) as br:
        with pytest.raises(ValidationError):
            await br.publish("wrong message", topic="test-topic")
            await handle.wait_call(timeout=3)

        handle.mock.assert_called_once_with("wrong message")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import pytest
from pydantic import ValidationError

from faststream.rabbit import TestRabbitBroker

from .pydantic_fields_rabbit import broker, handle


@pytest.mark.asyncio
async def test_handle():
    async with TestRabbitBroker(broker, with_real=True) as br:
        await br.publish({"name": "John", "user_id": 1}, queue="test-queue")
        await handle.wait_call(timeout=3)
        handle.mock.assert_called_once_with({"name": "John", "user_id": 1})

    assert handle.mock is None

@pytest.mark.asyncio
async def test_validation_error():
    async with TestRabbitBroker(broker, with_real=True) as br:
        with pytest.raises(ValidationError):
            await br.publish("wrong message", queue="test-queue")
            await handle.wait_call(timeout=3)

        handle.mock.assert_called_once_with("wrong message")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import pytest
from pydantic import ValidationError

from faststream.nats import TestNatsBroker

from .pydantic_fields_nats import broker, handle


@pytest.mark.asyncio
async def test_handle():
    async with TestNatsBroker(broker, with_real=True) as br:
        await br.publish({"name": "John", "user_id": 1}, subject="test-subject")
        await handle.wait_call(timeout=3)
        handle.mock.assert_called_once_with({"name": "John", "user_id": 1})

    assert handle.mock is None

@pytest.mark.asyncio
async def test_validation_error():
    async with TestNatsBroker(broker, with_real=True) as br:
        with pytest.raises(ValidationError):
            await br.publish("wrong message", subject="test-subject")
            await handle.wait_call(timeout=3)

        handle.mock.assert_called_once_with("wrong message")

Tip

When you're using a patched broker to test your consumers, the publish method is called synchronously with a consumer one, so you need not wait until your message is consumed. But in the real broker's case, it doesn't.

For this reason, you have to wait for message consumption manually with the special handler.wait_call(timeout) method. Also, inner handler exceptions will be raised in this function, not broker.publish(...).

A Little Tip#

It can be very helpful to set the with_real flag using an environment variable. This way, you will be able to choose the testing mode right from the command line:

WITH_REAL=True/False pytest tests/

To learn more about managing your application configiruation visit this page.

Last update: 2023-09-26