How to mock a database connection in FastAPI using polars?

Question

I want to mock pyodbc.Connection in a FastAPI application that uses the polars package to read the database. This is the initial main.py file:

from fastapi import Depends, FastAPI
import polars as pl
import pyodbc

import os

app = FastAPI()

def init_db() -> pyodbc.Connection:
    connstring = os.environ.get("AZURE_SQL_CONNECTIONSTRING", "no_env_var")
    return pyodbc.connect(connstring)


@app.get("/")
async def root(conn: pyodbc.Connection = Depends(init_db)) -> str:
    df = pl.read_database(query="SELECT * FROM test_table", connection=conn)
    return df.write_json()

The Depends class injects the database connection for testing afterward. I connected to the database via pl.read_database function. I want to create a unit test case without connecting to the real database.

How can I test the API calls without connecting to the real database using pytest?

EDIT: Replacing the to_dicts to write_json for more efficient conversion.

Answer 1

First of all, you can use the polars.ConnectionOrCursor class that can be returned when initializing the database:

from fastapi import Depends, FastAPI
import polars as pl
from polars.type_aliases import ConnectionOrCursor
import pyodbc

import os

app = FastAPI()

def init_db() -> ConnectionOrCursor:
    connstring = os.environ.get("AZURE_SQL_CONNECTIONSTRING", "no_env_var")
    return pyodbc.connect(connstring)


@app.get("/")
async def root(conn: ConnectionOrCursor = Depends(init_db)) -> str:
    df = pl.read_database(query="SELECT * FROM test_table", connection=conn)
    return df.write_json()

After that, it is possible to use a fake database in the unit test. You can use an in-memory SQLite database. Moreover, the TestClient class can be utilized for testing the FastAPI client. For more information visit the Testing FastAPI documentation.

This is an example test file:

from fastapi.testclient import TestClient
from polars.type_aliases import ConnectionOrCursor

import sqlite3

from your.source.module import app, init_db

persons = [
    (1, "John", 20),
    (2, "Doe", 30),
]


def setup_database() -> ConnectionOrCursor:
    connection = sqlite3.connect(":memory:", check_same_thread=False)
    cursor = connection.cursor()
    cursor.execute(
        "CREATE TABLE test_table(id, name, age)"
    )
    for person in persons:
        cursor.execute("INSERT INTO test_table VALUES(?,?,?,?)", person)
    connection.commit()

    return connection


client = TestClient(app)
app.dependency_overrides[init_db] = setup_database


def test_root():
    response = client.get("/")
    assert response.status_code == 200

If there are multiple test cases, use the check_same_thread=False in the connection constructor. It is possible to populate the SQLite database with data. Because the sqlite.Connection is compatible with Polars, you can override the init_db dependencies. The app.dependency_overrides attribute dictionary overrides the FastAPI dependencies.

Note: The table name (the test_table) has to be identical between the real and fake database.