Unit Test - pytest

YASE - Yet Another Software Engineer

Unit Test - pytest

To implement unit tests in Python, especially in our case with FastAPI, we use the TestClient object, which we can combine with the pytest package for more advanced functionalities.

Writing a test for an endpoint’s function boils down to these few steps:

# https://fastapi.tiangolo.com/tutorial/testing/#using-testclient
from fastapi import FastAPI
from fastapi.testclient import TestClient

app = FastAPI()

@app.get("/")
async def read_main():
    return {"msg": "Hello World"}

client = TestClient(app)

def test_read_main():
    payload = {"id": "test"} # Create payload
    response = client.post("/order", json=payload) # Make request
    assert response.status_code == 200 # Validate result
    assert response.json() == {"msg": "order created"}

For more complex cases involving the database, we can consider using mocks to isolate the test. This avoids having to delete results saved to the DB at the end of tests or having to create a database dedicated solely to testing.

A mock is nothing more than a simulated object that imitates the behavior of a real component (like a database, an external API, or another class).

Let’s see how to introduce this component into a test.

Mock

# https://fastapi.tiangolo.com/tutorial/testing/#using-testclient
import pytest # It's common to import pytest like this
from fastapi import FastAPI, Depends
from fastapi.testclient import TestClient

app = FastAPI()

def get_user():
    print("\nfoo") # Original function action
    return 'user-foo'

# Mocked function
def mock_get_user():
    print("\nmock-foo") # Mocked function action
    return 'user-mock'

@app.get("/")
async def read_main(user: str = Depends(get_user)):
    return {"msg": user}

@pytest.fixture(scope="function")  # Executed for each test function
def test_client_with_mock_user(): # More descriptive fixture name
    # Setup
    print("\n--- Fixture SETUP ---")
    original_get_user = app.dependency_overrides.get("get_user") # Store original if any
    app.dependency_overrides[get_user] = mock_get_user # Override by function object
    client = TestClient(app)
    yield client # Provide the TestClient
    # Teardown
    print("--- Fixture TEARDOWN ---")
    if original_get_user: # Restore original if it existed
        app.dependency_overrides[get_user] = original_get_user
    else:
        del app.dependency_overrides[get_user] # Or remove the override


def test_read_main_with_mock(test_client_with_mock_user: TestClient): # Use the fixture
    # No payload for GET
    response = test_client_with_mock_user.get("/")  # Make request
    print(response.text)
    assert response.status_code == 200  # Validate result
    assert response.json() == {"msg": "user-mock"}

The output will be:

============================== 1 passed in 0.30s ===============================

Process finished with exit code 0

--- Fixture SETUP ---
PASSED                                        [100%]
mock-foo
{"msg":"user-mock"}
--- Fixture TEARDOWN ---

The mock submodule is extremely powerful and flexible. One of its most common applications is to replace calls to the backend database. This serves two purposes: it ensures the test’s atomicity by not involving other components like the database, and it prevents unnecessary and potentially dangerous writes to the production database while tests are running.

Below is a trivial, almost useless, example of how to test an endpoint using a MagicMock object to replace the real database. Not only that, but after the request to the endpoint (which, thanks to the mock, is not harmful to the database even in case of errors), we can also verify if the endpoint’s handler actually performed all the operations we expected on the fake database object. This way, we can verify that the database interaction part also works without directly involving it in the test. Obviously, there are limitations; for instance, any DB component errors due to connection issues, users, permissions, or other factors will be difficult to detect using a mock.

# https://fastapi.tiangolo.com/tutorial/testing/#using-testclient
from unittest.mock import MagicMock

import pytest # Standard import
from fastapi import FastAPI, Depends
from fastapi.testclient import TestClient
from sqlmodel import select, SQLModel, Field, Session # Assuming SQLModel

app = FastAPI()
# db_mock = MagicMock(spec=Session) # It's good practice to spec the MagicMock

# Define a User model for context
class User(SQLModel, table=True):
    id: int | None = Field(default=None, primary_key=True)
    name: str
    password: str # In a real app, this would be a hash
    age: int | None = None

# This would be your actual dependency in the app
def get_db_real():
    print("official get_db function called - THIS SHOULD NOT HAPPEN IN TEST")
    # Real database session logic here
    raise NotImplementedError("Real DB access in test")


# This is the dependency we will override
@app.get("/")
async def read_user_from_db(db: Session = Depends(get_db_real)): # Depends on the real one initially
    query = select(User).where(User.name == "spongebob")
    # In a real scenario, you'd expect db.exec(query).all() or .first()
    # For MagicMock, we often mock the return value of chained calls
    # For simplicity, let's assume db.execute returns an object that has an all() method
    # db.execute(query).all()
    db.execute(query) # Let's say our app code just calls execute for now
    # And then maybe it processes results... for this test, we'll just check execute was called.
    # This endpoint as written doesn't actually return the result of db.execute,
    # which is a bit unusual. Let's assume it does something with it internally or returns a fixed value.
    return {"status": "query executed"}


@pytest.fixture(scope="function")
def client_with_mock_db():
    # Setup
    print("\n--- Fixture SETUP for DB Mock ---")
    # Create a fresh mock for each test to avoid state leakage
    mock_db_session = MagicMock(spec=Session)
    
    # If your code does something like: result = db.exec(query).all(), you might mock it like this:
    # mock_result_proxy = MagicMock()
    # mock_result_proxy.all.return_value = [User(id=1, name="spongebob", password="pwd")] # Example return
    # mock_db_session.exec.return_value = mock_result_proxy

    app.dependency_overrides[get_db_real] = lambda: mock_db_session # Override with a function returning the mock
    
    test_client = TestClient(app)
    yield test_client, mock_db_session # Yield both client and mock for assertions

    # Teardown
    print("--- Fixture TEARDOWN for DB Mock ---")
    app.dependency_overrides.clear() # Clear all overrides


def test_read_main_with_db_mock(client_with_mock_db: tuple[TestClient, MagicMock]):
    test_client, mock_db = client_with_mock_db
    # No payload for GET requests
    response = test_client.get("/")  # Make request
    
    # [https://docs.python.org/3/library/unittest.mock.html#unittest.mock.Mock.assert_called_once](https://docs.python.org/3/library/unittest.mock.html#unittest.mock.Mock.assert_called_once)
    # Verify that the execute method was called at most once
    # This is a basic test, but it's just an example of the mock module's capabilities
    mock_db.execute.assert_called_once()  # Validate interaction with the mock
    
    # You could also check what it was called with:
    # from unittest.mock import ANY (or capture_first_arg from a helper)
    # mock_db.execute.assert_called_once_with(ANY) # Or a more specific query object if constructible

    assert response.status_code == 200
    # Adjust assertion based on what the endpoint actually returns
    assert response.json() == {"status": "query executed"}

In the code above, albeit trivial, we’ve verified that an “execute” was actually performed on the DB and that this function was called once and only once. We could have written a more useful test, perhaps by checking that a specific query was executed, its parameters, etc., but this is just an example. In the code, you’ll find the link to the mock documentation to explore all the verification functionalities associated with MagicMock.