FastAPI+BigQuery后端单元测试问题：如何Mock BigQuery Client

首先，你的代码存在一个关键问题：在get_price端点中，你忽略了通过Depends注入的bq_client，反而重新创建了新的客户端实例，这会导致依赖注入失效，无法通过替换依赖来Mock BigQuery服务。先修正这一问题：

修正main.py中的端点代码 #

from google.cloud import bigquery
from fastapi import FastAPI, Depends
import pandas as pd

app = FastAPI()

def get_bq_client():
    # BigQuery客户端本身线程安全，无需上下文管理器包裹，直接返回实例即可
    return bigquery.Client()

def prepare_query(val1):
    query = f"""
         SELECT name, item_id, price
         FROM `bigquery-my-data.products`
         WHERE item_id = {val1}
     """
    return query

@app.get("/v1/items/{item_id}")
async def get_price(item_id: int,
                    bq_client: bigquery.client.Client = Depends(get_bq_client)
                    ):
    query = prepare_query(item_id)
    # 使用注入的客户端实例，而非重新创建
    query_job = bq_client.query(query)
    # 原代码中query_job.query()是错误调用，应使用result()获取查询结果
    result = query_job.result().to_dataframe()
    
    return result.to_json()

编写单元测试（Mock BigQuery客户端） #

使用Python标准库unittest.mock来模拟BigQuery客户端及相关方法，避免实际调用云端服务：

from fastapi.testclient import TestClient
from unittest.mock import Mock, patch
from main import app
import pandas as pd
import json

client = TestClient(app)

def test_get_price_correct():
    # 模拟BigQuery返回的DataFrame数据
    mock_df = pd.DataFrame({
        "name": ["Cola"],
        "item_id": [123],
        "price": [1.12]
    })
    
    # 逐层Mock：QueryJob -> result() -> to_dataframe()
    mock_query_job = Mock()
    mock_query_job.result.return_value.to_dataframe.return_value = mock_df
    
    # Mock BigQuery客户端类及其query方法
    with patch("main.bigquery.Client") as mock_client_cls:
        mock_client_instance = Mock()
        mock_client_instance.query.return_value = mock_query_job
        mock_client_cls.return_value = mock_client_instance
        
        # 发送测试请求
        response = client.get("/v1/items/123")
        
        # 验证响应状态与数据
        assert response.status_code == 200
        response_data = json.loads(response.json())
        expected = {
            "name": {"0": "Cola"},
            "item_id": {"0": 123},
            "price": {"0": 1.12}
        }
        assert response_data == expected
        
        # 验证查询语句是否正确生成
        mock_client_instance.query.assert_called_once()
        called_query = mock_client_instance.query.call_args[0][0]
        assert "WHERE item_id = 123" in called_query

关键要点说明 #

• 依赖注入修正：确保端点使用Depends注入的客户端，才能在测试中替换为Mock实例。
• Mock层级：需依次模拟bigquery.Client实例、query方法返回的QueryJob、result()方法，最终模拟to_dataframe()返回的测试数据。
• 验证逻辑：不仅要检查响应数据，还要验证查询语句是否正确生成，确保业务逻辑符合预期。
• 安全优化建议：原prepare_query使用字符串拼接易引发SQL注入，建议改为参数化查询：

  def prepare_query(val1):
      query = """
           SELECT name, item_id, price
           FROM `bigquery-my-data.products`
           WHERE item_id = @val1
       """
      job_config = bigquery.QueryJobConfig(
          query_parameters=[
              bigquery.ScalarQueryParameter("val1", "INT64", val1)
          ]
      )
      return query, job_config
  

  端点中对应修改为：

  query, job_config = prepare_query(item_id)
  query_job = bq_client.query(query, job_config=job_config)
  

内容的提问来源于stack exchange，提问作者Roberto