使用 Tortoise-ORM 在 FastAPI 中进行测试

Question

我正在尝试在 Python 3.8 下使用 Tortoise ORM 在 FastAPI 中编写一些异步测试，但我不断收到相同的错误（见最后）。过去几天我一直在尝试解决这个问题，但不知何故，我最近创建测试的所有努力都没有成功。

我正在关注这个fastapi docs 和tortoise docs。

ma​​in.py

# UserPy is a pydantic model
@app.post('/testpost')
async def world(user: UserPy) -> UserPy:
    await User.create(**user.dict())
    # Just returns the user model
    return user

simple_test.py

from fastapi.testclient import TestClient
from httpx import AsyncClient

@pytest.fixture
def client1():
    with TestClient(app) as tc:
        yield tc

@pytest.fixture
def client2():
    initializer(DATABASE_MODELS, DATABASE_URL)
    with TestClient(app) as tc:
        yield tc
    finalizer()

@pytest.fixture
def event_loop(client2):              # Been using client1 and client2 on this
    yield client2.task.get_loop()


# The test
@pytest.mark.asyncio
def test_testpost(client2, event_loop):
    name, age = ['sam', 99]
    data = json.dumps(dict(username=name, age=age))
    res = client2.post('/testpost', data=data)
    assert res.status_code == 200

    # Sample query
    async def getx(id):
        return await User.get(pk=id)
    x = event_loop.run_until_complete(getx(123))
    assert x.id == 123

    # end of code

我的错误取决于我使用的是client1 还是client2

使用client1 错误

RuntimeError: Task <Task pending name='Task-9' coro=<TestClient.wait_shutdown() running at <my virtualenv path>/site-packages/starlette/testclient.py:487> cb=[_run_until_complete_cb() at /usr/lib/python3.8/asyncio/base_events.py:184]> got Future <Future pending> attached to a different loop

使用client2 错误

asyncpg.exceptions.ObjectInUseError: cannot drop the currently open database

哦，我也尝试过使用httpx.AsyncClient，但仍然没有成功（还有更多错误）。任何想法，因为我没有自己的想法。

Answer 1

我花了大约一小时来完成异步测试。这是示例： （需要Python3.8+）

• conftest.py

import asyncio

import pytest
from tortoise import Tortoise

DB_URL = "sqlite://:memory:"


async def init_db(db_url, create_db: bool = False, schemas: bool = False) -> None:
    """Initial database connection"""
    await Tortoise.init(
        db_url=db_url, modules={"models": ["models"]}, _create_db=create_db
    )
    if create_db:
        print(f"Database created! {db_url = }")
    if schemas:
        await Tortoise.generate_schemas()
        print("Success to generate schemas")


async def init(db_url: str = DB_URL):
    await init_db(db_url, True, True)


@pytest.fixture(scope="session")
def event_loop():
    return asyncio.get_event_loop()


@pytest.fixture(scope="session", autouse=True)
async def initialize_tests():
    await init()
    yield
    await Tortoise._drop_databases()

• settings.py

import os

from dotenv import load_dotenv

load_dotenv()

DB_NAME = "async_test"
DB_URL = os.getenv(
    "APP_DB_URL", f"postgres://postgres:postgres@127.0.0.1:5432/{DB_NAME}"
)

ALLOW_ORIGINS = [
    "http://localhost",
    "http://localhost:8080",
    "http://localhost:8000",
    "https://example.com",
]

• main.py

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from models.users import User, User_Pydantic, User_Pydantic_List, UserIn_Pydantic
from settings import ALLOW_ORIGINS, DB_URL
from tortoise.contrib.fastapi import register_tortoise

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOW_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.post("/testpost", response_model=User_Pydantic)
async def world(user: UserIn_Pydantic):
    return await User.create(**user.dict())


@app.get("/users", response_model=User_Pydantic_List)
async def user_list():
    return await User.all()


register_tortoise(
    app,
    config={
        "connections": {"default": DB_URL},
        "apps": {"models": {"models": ["models"]}},
        "use_tz": True,
        "timezone": "Asia/Shanghai",
        "generate_schemas": True,
    },
)

• models/base.py

from typing import List, Set, Tuple, Union

from tortoise import fields, models
from tortoise.query_utils import Q
from tortoise.queryset import QuerySet


def reduce_query_filters(args: Tuple[Q, ...]) -> Set:
    fields = set()
    for q in args:
        fields |= set(q.filters)
        c: Union[List[Q], Tuple[Q, ...]] = q.children
        while c:
            _c: List[Q] = []
            for i in c:
                fields |= set(i.filters)
                _c += list(i.children)
            c = _c
    return fields


class AbsModel(models.Model):
    id = fields.IntField(pk=True)
    created_at = fields.DatetimeField(auto_now_add=True, description="Created At")
    updated_at = fields.DatetimeField(auto_now=True, description="Updated At")
    is_deleted = fields.BooleanField(default=False, description="Mark as Deleted")

    class Meta:
        abstract = True
        ordering = ("-id",)

    @classmethod
    def filter(cls, *args, **kwargs) -> QuerySet:
        field = "is_deleted"
        if not args or (field not in reduce_query_filters(args)):
            kwargs.setdefault(field, False)
        return super().filter(*args, **kwargs)

    class PydanticMeta:
        exclude = ("created_at", "updated_at", "is_deleted")

    def __repr__(self):
        return f"<{self.__class__.__name__} {self.id}>"

• models/users.py

from tortoise.contrib.pydantic import pydantic_model_creator, pydantic_queryset_creator

from .base import AbsModel, fields


class User(AbsModel):
    username = fields.CharField(60)
    age = fields.IntField()

    class Meta:
        table = "users"

    def __str__(self):
        return self.name


User_Pydantic = pydantic_model_creator(User)
UserIn_Pydantic = pydantic_model_creator(User, name="UserIn", exclude_readonly=True)
User_Pydantic_List = pydantic_queryset_creator(User)

• models/__init__.py

from .users import User  # NOQA: F401

• 测试/test_users.py

import pytest
from httpx import AsyncClient
from main import app
from models.users import User


@pytest.mark.asyncio
async def test_testpost():
    name, age = ["sam", 99]
    assert await User.filter(username=name).count() == 0

    data = {"username": name, "age": age}
    async with AsyncClient(app=app, base_url="http://test") as ac:
        response = await ac.post("/testpost", json=data)
        assert response.json() == dict(data, id=1)
        assert response.status_code == 200

        response = await ac.get("/users")
        assert response.status_code == 200
        assert response.json() == [dict(data, id=1)]

    assert await User.filter(username=name).count() == 1

演示的源代码已发布到github： https://github.com/waketzheng/fastapi-tortoise-pytest-demo.git