PsycologyAPI/main.py

from contextlib import asynccontextmanager
from random import randint
from typing import AsyncGenerator

from fastapi import FastAPI
from tortoise import Tortoise, generate_config
from tortoise.contrib.fastapi import RegisterTortoise

from models import RawData


@asynccontextmanager
async def lifespan_test(app: FastAPI) -> AsyncGenerator[None, None]:
    config = generate_config(
        "sqlite://:memory:",
        app_modules={"models": ["models"]},
        testing=True,
        connection_label="models",
    )
    async with RegisterTortoise(
            app=app,
            config=config,
            generate_schemas=True,
            add_exception_handlers=True,
            _create_db=True,
    ):
        # db connected
        yield
        # app teardown
    # db connections closed
    await Tortoise._drop_databases()


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
    config = generate_config(
        "sqlite://db.sqlite3",
        app_modules={"models": ["models"]},
        connection_label="models",
    )
    async with RegisterTortoise(
            app=app,
            config=config
    ):
        # db connected
        yield
        # app teardown
    # db connections closed
    await Tortoise.close_connections()


app = FastAPI(lifespan=lifespan)


@app.get("/")
async def root():
    cnt = await RawData.all().count()
    # 生成一个范围内的随机整数，随机返回一篇文章
    r = randint(1, cnt)
    record = await RawData.get(id=r)
    return {"title": record.title, "content": record.content}