如何在使用Polars的FastAPI应用中Mock数据库连接并编写pytest单元测试？

没问题！要在不连接真实数据库的前提下测试你的FastAPI接口，我们可以借助pytest和Python自带的unittest.mock工具，结合FastAPI的依赖注入特性来实现。下面是具体的步骤和代码示例，两种思路任你选：

1. 先装好用得上的测试依赖 #

首先确保你安装了测试所需的包，打开终端执行：

pip install pytest httpx

• pytest是我们的测试框架，httpx用来模拟HTTP请求调用FastAPI接口。

2. 编写测试文件（建议命名为test_main.py） #

思路一：直接模拟Polars的read_database方法 #

这种方式不用纠结数据库连接的细节，直接让Polars返回我们预设的测试数据，适合快速验证接口的返回逻辑：

from fastapi.testclient import TestClient
import polars as pl
from unittest.mock import patch
from main import app

# 初始化测试客户端
client = TestClient(app)

def test_root_endpoint_mock_polars():
    # 先准备好测试用的模拟数据
    test_df = pl.DataFrame({
        "id": [1, 2, 3],
        "name": ["Alice", "Bob", "Charlie"]
    })
    
    # 用patch模拟polars.read_database方法，让它返回我们的测试DataFrame
    with patch("polars.read_database") as mock_read_db:
        mock_read_db.return_value = test_df
        
        # 发送GET请求到接口根路径
        response = client.get("/")
        
        # 验证请求是否成功（状态码200）
        assert response.status_code == 200
        
        # 验证返回的JSON和我们预设的DataFrame转换后的JSON一致
        assert response.json() == test_df.to_dicts()
        
        # 还可以验证read_database是否被正确调用了（检查查询语句和连接参数）
        mock_read_db.assert_called_once_with(
            query="SELECT * FROM test_table",
            connection=mock_read_db.call_args[1]["connection"]
        )

思路二：模拟依赖注入的数据库连接 #

这种方式更贴合你代码里的依赖注入设计，替换掉真实的init_db函数，返回一个模拟的pyodbc.Connection对象，确保Polars正确使用了这个连接：

from fastapi.testclient import TestClient
import polars as pl
import pyodbc
from unittest.mock import Mock, patch
from main import app, init_db

client = TestClient(app)

def test_root_endpoint_mock_db_connection():
    # 同样先准备测试数据
    test_df = pl.DataFrame({
        "id": [1, 2, 3],
        "name": ["Alice", "Bob", "Charlie"]
    })
    
    # 创建一个模拟的pyodbc.Connection对象（用spec参数保证它符合真实连接的接口）
    mock_conn = Mock(spec=pyodbc.Connection)
    
    # 定义一个替换init_db的函数，让它返回我们的模拟连接
    def mock_init_db():
        return mock_conn
    
    # 替换main模块里的init_db函数
    with patch("main.init_db", side_effect=mock_init_db):
        # 同时模拟polars.read_database返回测试数据
        with patch("polars.read_database") as mock_read_db:
            mock_read_db.return_value = test_df
            
            # 发送请求
            response = client.get("/")
            
            # 验证请求状态
            assert response.status_code == 200
            # 验证返回内容正确性
            assert response.json() == test_df.to_dicts()
            
            # 验证read_database确实使用了我们模拟的连接
            mock_read_db.assert_called_once_with(
                query="SELECT * FROM test_table",
                connection=mock_conn
            )

3. 运行测试 #

在终端里执行下面的命令，就能看到测试结果了：

pytest test_main.py -v

额外的小提示 #

如果你还想单独测试init_db函数的逻辑（比如验证它是否正确读取环境变量），可以加个单独的测试用例：

import os
from unittest.mock import patch
from main import init_db

def test_init_db_env_var():
    # 临时设置环境变量
    with patch.dict(os.environ, {"AZURE_SQL_CONNECTIONSTRING": "test_conn_string"}):
        with patch("pyodbc.connect") as mock_connect:
            init_db()
            # 验证pyodbc.connect是否用了环境变量里的连接字符串
            mock_connect.assert_called_once_with("test_conn_string")

def test_init_db_default_value():
    # 清空环境变量的情况
    with patch.dict(os.environ, {}, clear=True):
        with patch("pyodbc.connect") as mock_connect:
            init_db()
            # 验证用了默认的"no_env_var"连接字符串
            mock_connect.assert_called_once_with("no_env_var")

备注：内容来源于stack exchange，提问作者Péter Szilvási