Managing Exceptions in WebSockets with FastAPI

In this post, we delve into the management of exceptions in WebSockets, focusing on a potent but often overlooked feature: the WebSocketException offered by Starlette.

Understanding WebSocketException

Conceptually, WebSocketException enables you to close a WebSocket connection with a specific code and reason by raising this exception in your WebSocket route. Here's an illustrative example:

from fastapi import FastAPI, WebSocket
from fastapi.exceptions import WebSocketException

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    raise WebSocketException(code=1008, reason="Closing the connection...")

In this instance, a WebSocketException is raised bearing the code 1008 and the explicit reason for closure: "Closing the connection...".

To run this application, first install FastAPI, Uvicorn, and WebSockets:

pip install fastapi uvicorn websockets

Run the application using Uvicorn:

uvicorn main:app --reload

On opening a WebSocket connection to ws://localhost:8000/ws, you will find the connection being closed with focal code 1008 and the attributed reason.

I use wscat for testing WebSocket connections; you can install it with the following command:

npm install -g wscat

A connection is opened with:

wscat -c ws://127.0.0.1:8000/ws

Handle custom exceptions

The app.exception_handler can be used to handle custom exceptions. Consider the following sample:

from fastapi import FastAPI, WebSocket
from fastapi.exceptions import WebSocketException

app = FastAPI()

class CustomException(Exception): ...

@app.exception_handler(CustomException)
async def custom_exception_handler(websocket: WebSocket, exc: Exception):
    await websocket.close(code=1008, reason="Custom exception")

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    raise CustomException()

In this example, the client receives a WebSocket closure with code 1008 and the reason "Custom exception" when the CustomException is raised.

Feel free to connect with me on LinkedIn for any questions or discussions on this topic.

Happy coding! 🚀

Contract Testing with HTTPX - Part 2

Note

This is a continuation of Contract Testing with HTTPX.

On the previous article, I used RESPX to call the service B from service A. Although it looks cool, we can actually achieve the same goal without using anything besides FastAPI itself.

The Services

Let's assume that we have similar two services as presented in the previous article.

The difference here is that we'll be creating a dependency called service_b_client, which is going to return a httpx.AsyncClient instance that calls our service B.

service_a.py
from typing import AsyncIterator

import httpx
from fastapi import APIRouter, Depends, FastAPI

router = APIRouter(prefix="/a")


async def service_b_client() -> AsyncIterator[httpx.AsyncClient]:
    async with httpx.AsyncClient(base_url="http://localhost:8000/b/") as client:
        yield client


@router.get("/")
def get_a():
    return {"a": "a"}


@router.get("/call_b")
async def call_b(client: httpx.AsyncClient = Depends(service_b_client)):
    response = await client.get("/")
    return response.json()


app = FastAPI()
app.include_router(router)

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, port=8001)  # (1)!
  1. The port is 8001, not 8000, to avoid conflicts with the other service.
service_b.py
from fastapi import APIRouter, FastAPI

router = APIRouter(prefix="/b")


@router.get("/")
def get_b():
    return {"b": "b"}


app = FastAPI()
app.include_router(router)

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, port=8000)

Install the dependencies:

python -m pip install uvicorn fastapi httpx

Then open the terminal and run:

python service_a.py

Then open another terminal, and run:

python service_b.py

Now, let's call the /a/call_b endpoint.

http :8001/a/call_b # (1)!
  1. The HTTP client used is called HTTPie, but you can use [curl], or just go to the browser, and access http://localhost:8001/a/call_b.

The response should look like:

{
    "b": "b"
}

Testing

Since the only difference between this article, and the previous one is the creation of the dependency on service A, you might be guessing that we are going to override the dependency, and... You are right! (if you didn't, is fine as well 😅)

We are going to use app.dependency_overrides to override the service_b_client dependency, and instead of calling the real service B, we'll call the application itself, avoiding the network calls that would potentially slow down our test suite.

test.py
from typing import AsyncIterator

import httpx
import pytest

from service_a import app, service_b_client
from service_b import app as app_b


async def service_b_client_override() -> AsyncIterator[httpx.AsyncClient]:
    async with httpx.AsyncClient(app=app_b, base_url="http://test/b") as client:
        yield client


@pytest.fixture(name="client")
async def testclient():
    app.dependency_overrides[service_b_client] = service_b_client_override
    async with httpx.AsyncClient(app=app, base_url="http://test") as client:
        yield client


@pytest.mark.anyio
async def test_call_b(client: httpx.AsyncClient) -> None:
    response = await client.get("/a/call_b")
    assert response.status_code == 200
    assert response.json() == {"b": "b"}

See more on the Starlette documentation.

Install the dependencies:

python -m pip install pytest httpx trio

Then open the terminal and run:

python -m pytest test.py

Good! Now are able to avoid a lot of network calls, and speed up our test suite. 🎉

Contract Testing with HTTPX

Today, we are going to talk about how to achieve contract testing with HTTPX. 🤓

What is contract testing?

Contract testing is a methodology for ensuring that two separate systems (such as two microservices) are compatible and are able to communicate with one other. It captures the interactions that are exchanged between each service, storing them in a contract, which can then be used to verify that both parties adhere to it. - Matt Fellows

What is HTTPX?

HTTPX is a fully featured HTTP client for Python 3, which provides sync and async APIs, and support for both HTTP/1.1 and HTTP/2.

How to do contract testing with HTTPX?

Well, to be completely transparent, I'm not sure if what you are about to read classifies as contract testing. 😅

The problem we'll be trying to solve is the following:

Consider we have multiples services running, and they depend on each other. We want to make sure that a service is not able able to break another one.

To achieve this your first thought would be "let's write end to end tests", but that will slow things down, as each service needs to be up to run the tests, and given that, the setup needed is a bit more complex.

Check this blog post (which I didn't read, but looks good) for more information about E2E testing vs Contract Testing.

The solution

Let's assume we have two services. For obvious reasons, those services are FastAPI based. 👀

Note

This can be achieved with any web framework. What matters here is that you should be using httpx.

service_a.py
import httpx
from fastapi import APIRouter, FastAPI

router = APIRouter(prefix="/a")


@router.get("/")
def get_a():
    return {"a": "a"}


@router.get("/call_b")
async def call_b():
    async with httpx.AsyncClient() as client:
        response = await client.get("http://localhost:8000/b/")
        return response.json()


app = FastAPI()
app.include_router(router)

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, port=8001)  # (1)!
  1. The port is 8001, not 8000, to avoid conflicts with the other service.
service_b.py
from fastapi import APIRouter, FastAPI

router = APIRouter(prefix="/b")


@router.get("/")
def get_b():
    return {"b": "b"}


app = FastAPI()
app.include_router(router)

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, port=8000)

Install the dependencies:

python -m pip install uvicorn fastapi httpx

Then open the terminal and run:

python service_a.py

Then open another terminal, and run:

python service_b.py

Cool. Now, let's call the /a/call_b endpoint:

http :8001/a/call_b # (1)!
  1. The HTTP client used is called HTTPie, but you can use curl, or just go to the browser, and access http://localhost:8001/a/call_b.

As we see, the response is:

{
    "b": "b"
}

Now, if we want to create a test, we can do something like:

test_service_a.py
import pytest
from fastapi.testclient import TestClient

from service_a import app


@pytest.fixture()
def client():
    return TestClient(app)


def test_call_b(client: TestClient) -> None:
    response = client.get("/a/call_b")
    assert response.status_code == 200
    assert response.json() == {"b": "b"}

See more on the Starlette documentation.

test_service_a.py
import httpx
import pytest
import pytest_asyncio

from service_a import app


@pytest_asyncio.fixture(name="client")
async def testclient():
    async with httpx.AsyncClient(app=app, base_url="http://test") as client:
        yield client


@pytest.mark.asyncio
async def test_call_b(client: httpx.AsyncClient) -> None:
    response = await client.get("/a/call_b")
    assert response.status_code == 200
    assert response.json() == {"b": "b"}

See more on the Starlette documentation.

Install the dependencies:

python -m pip install pytest pytest-asyncio httpx requests

Then open the terminal and run:

python -m pytest test_service_a.py

That works perfectly, right? 🤔

Well, yes. But remember what I said in the beginning? 😅

To achieve this your first thought would be "let's write end to end tests", but that will slow things down, as each service needs to be up to run the tests, and given that, the setup needed is a bit more complex.

So, what if we want to run the tests without having to run the services? 🤔

Patch the HTTP client

We can patch the HTTPX client to make it call the service B, without actually running the service B. ✨

To achieve that, we'll be using RESPX: a simple library, yet powerful, for mocking out the HTTPX and HTTPCore libraries. ✨

It's easy, let me show you! We just need to add a single fixture on the test_service_a.py file:

test_service_a.py
import pytest
import pytest_asyncio
import respx
from fastapi.testclient import TestClient

from service_a import app
from service_b import app as app_b


@pytest.fixture()
def client():
    return TestClient(app)


@pytest_asyncio.fixture()
async def create_contracts():
    async with respx.mock:  # (1)!
        respx.route(host="localhost", port=8000).mock(
            side_effect=respx.ASGIHandler(app_b)
        )
        yield


def test_call_b(client: TestClient) -> None:
    response = client.get("/a/call_b")
    assert response.status_code == 200
    assert response.json() == {"b": "b"}
  1. The respx.mock context manager is used to mock the HTTPX client. ✨

    Read more about it on the RESPX documentation.

test_service_a.py
import httpx
import pytest
import pytest_asyncio
import respx

from service_a import app
from service_b import app as app_b


@pytest_asyncio.fixture(name="client")
async def testclient():
    async with httpx.AsyncClient(app=app, base_url="http://test") as client:
        yield client


@pytest_asyncio.fixture()
async def create_contracts():
    async with respx.mock:  # (1)!
        respx.route(host="localhost", port=8000).mock(
            side_effect=respx.ASGIHandler(app_b)
        )
        yield


@pytest.mark.asyncio
async def test_call_b(client: httpx.AsyncClient) -> None:
    response = await client.get("/a/call_b")
    assert response.status_code == 200
    assert response.json() == {"b": "b"}
  1. The respx.mock context manager is used to mock the HTTPX client. ✨

    Read more about it on the RESPX documentation.

Install the dependencies:

python -m pip install pytest pytest-asyncio httpx respx requests

Then open the terminal and run:

python -m pytest test_service_a.py

Nice! We did it! 🎉

Now, we can run the tests without having to run the services. 🚀

If you are a curious person, feel free to compare the tests with the time command:

time python -m pytest test_service_a.py

Be surprised. 😉

Info

You can also read the continuation of this article here.

FastAPI Escape Character

Today, we'll talk about a small feature of FastAPI that might be useful for you: the escape character. 🤓

What is the escape character?

The escape character \f is a character that can be used to tell to FastAPI to truncate what should go to the endpoint description on the OpenAPI.

Let's see it in practice. Consider we have the following endpoint:

main.py
from fastapi import FastAPI

app = FastAPI()

@app.get("/")
def home():
    """This is home.
    \f
    This is not on the OpenAPI.
    """

Install the dependencies:

python -m pip install uvicorn fastapi

Then run uvicorn:

uvicorn main:app

When we call the /openapi.json endpoint:

http GET :8000/openapi.json  # (1)!
  1. The HTTP client used is called HTTPie, but you can use curl, or just go to the browser, and access http://localhost:8000/openapi.json.

You'll see the following OpenAPI JSON on the response:

{
    "info": {
        "title": "FastAPI",
        "version": "0.1.0"
    },
    "openapi": "3.0.2",
    "paths": {
        "/": {
            "get": {
                "description": "This is home.",
                "operationId": "home__get",
                "responses": {
                    "200": {
                        "content": {
                            "application/json": {
                                "schema": {}
                            }
                        },
                        "description": "Successful Response"
                    }
                },
                "summary": "Home"
            }
        }
    }
}

Observe the "description" field does not contain the "This is not on OpenAPI" part of the docstring. The reason is the escape character we used. Everything after the \f will not appear on that field.

This feature may be useful if you are using a docstring linter tool, like darglint.

What's new?

If you are a FastAPI veteran (😎), you are probably familiar with the above. What you probably don't know, is that now (since FastAPI 0.82.0) it's possible to use it on the Pydantic models you use on your FastAPI application.

Let's see another example.

As most of my examples, we'll use potatoes:

main.py
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

class PotatoOutput(BaseModel):
    """Super potato.
    \f
    This is not on the OpenAPI.
    """

@app.get("/", response_model=PotatoOutput)
def get_potato():
    ...

Install the dependencies:

python -m pip install uvicorn fastapi

Then run uvicorn:

uvicorn main:app

When we call /openapi.json, as we did above, we'll get the following OpenAPI JSON as response:

{
    "components": {
        "schemas": {
            "PotatoOutput": {
                "description": "Super potato.\n",
                "properties": {},
                "title": "PotatoOutput",
                "type": "object"
            }
        }
    },
    "info": {
        "title": "FastAPI",
        "version": "0.1.0"
    },
    "openapi": "3.0.2",
    "paths": {
        "/": {
            "get": {
                "operationId": "get_potato__get",
                "responses": {
                    "200": {
                        "content": {
                            "application/json": {
                                "schema": {
                                    "$ref": "#/components/schemas/PotatoOutput"
                                }
                            }
                        },
                        "description": "Successful Response"
                    }
                },
                "summary": "Get Potato"
            }
        }
    }
}

Tip

We can also use jq to get the part of the JSON that we are interested.

http GET :8000/openapi.json | jq .components.schemas
{
    "PotatoOutput": {
        "title": "PotatoOutput",
        "type": "object",
        "properties": {},
        "description": "Super potato.\n"
    }
}

As we can see, the description of PotatoOutput doesn't contain the "This is not on the OpenAPI." part as well.

Yey! Now you can use those docstring linter tools as you want with FastAPI! 🙌


Thanks for reading this blog post! 🥳

If you have any suggestions on what I can write about, please feel free to suggest below. 🙏

FastAPI's Test Client

This is my first blog post! 🥳

Please enjoy, and let me know if you have any feedback. 🤓

Abstract

If you are new to FastAPI, you might benefit from reading the following:

If you already know stuff about FastAPI, you might jump to:

Today, we'll talk about the main tool for testing FastAPI applications: the TestClient.

TestClient origin and features

The TestClient is a feature from Starlette (one of the two main dependencies of FastAPI). On which, FastAPI only does a reimport on the testclient module, as we can see here.

We can use the TestClient to test our WebSocket and HTTP endpoints.

The TestClient weird behavior

Although documented on both FastAPI and Starlette's documentation, most of the people are not aware of the TestClient's behavior when it comes to events. To put it simple, there are two ways of creating a TestClient object, and in one of those ways, the events are not executed.

Let's see the behavior with the following FastAPI application:

main.py
from fastapi import FastAPI

app = FastAPI()
started = False

@app.on_event("startup") # (1)!
def startup():
    global started
    started = True

@app.get("/")
def home():
    if started:
        return {"message": "STARTED"}
    else:
        return {"message": "NOT STARTED"}
  1. There are only two events available: startup and shutdown.

    Read more about it on the ASGI documentation.

As you can see, there's a single endpoint, which gives us a different message depending on the value of the started variable. The started variable is set to True on the startup event.

Now, let's test it with the TestClient:

test.py
1
2
3
4
5
6
7
8
9
from fastapi.testclient import TestClient

from main import app

def test_home():
    client = TestClient(app)
    response = client.get("/")
    assert response.status_code == 200
    assert response.json() == {"message": "NOT STARTED"}

Install the dependencies:

python -m pip install "fastapi[all]" pytest

Then run pytest:

pytest test.py


As you can see, the above test passes. Which means the startup event was not triggered. 😱

On the other hand, if we run the following test, we'll get a different result:

test.py
1
2
3
4
5
6
7
8
9
from fastapi.testclient import TestClient

from main import app

def test_home():
    with TestClient(app):
        response = client.get("/")
        assert response.status_code == 200
        assert response.json() == {"message": "STARTED"}

Install the dependencies:

python -m pip install "fastapi[all]" pytest

Then run pytest:

pytest test.py


When used as context manager, the TestClient will trigger the startup event.

The Future of the TestClient

By the moment I'm writing this blog, the latest FastAPI version is 0.83.0 with Starlette pinned on 0.19.1. Starlette is already on version 0.20.3, and the next release will change the internals of the TestClient. To be more specific, the HTTP client will be changed from requests to httpx.

As there are some differences between the two clients, the TestClient will reflect the same differences.

This change will be in Starlette on version 0.21.0, and I'm unsure when it will land on FastAPI.

Let's see the changes you should be aware:

  1. allow_redirects will be now called follow_redirects.
  2. cookies parameter will be deprecated under method calls (it should be used on the client instantiation).
  3. data parameter will be called content when sending bytes or text.
  4. content_type will default to "text/plain" when sending file instead of empty string.
  5. The HTTP methods DELETE, GET, HEAD and OPTIONS will not accept content, data, json and files parameters.
  6. data parameter doesn't accept list of tuples, instead it should be a dictionary.

    client.post(..., data=[("key1", "1"), ("key1", "2"), ("key2", "3")])
    
    client.post(..., data={"key1": ["1", "2"], "key2": "3"})
    

Those changes will likely impact your test suite. Having this in mind, I've created a codemod that will help you to migrate your tests: bump-testclient. 🎉

Here is the list of what the codemod will do:

  1. Replace allow_redirects with follow_redirects.
  2. Replace data with content when sending bytes or text.
  3. Replace client.<method>(..., <parameter>=...) by client.request("<method>", ..., <parameter>=...) when parameter is either content, data, json or files.

In case you want to read more about the differences between the underneath clients, you can check the httpx documentation.

Thanks for reading till here! 🤓