from beanie.odm.queries.update import UpdateOne
from tests.typing.models import ProjectionTest, Test


async def find_many() -> list[Test]:
    return await Test.find().to_list()


async def find_many_with_projection() -> list[ProjectionTest]:
    return await Test.find().project(projection_model=ProjectionTest).to_list()


async def find_many_generator() -> list[Test]:
    docs: list[Test] = []
    async for doc in Test.find():
        docs.append(doc)
    return docs


async def find_many_generator_with_projection() -> list[ProjectionTest]:
    docs: list[ProjectionTest] = []
    async for doc in Test.find().project(projection_model=ProjectionTest):
        docs.append(doc)
    return docs


async def find_one() -> Test | None:
    return await Test.find_one()


async def find_one_with_projection() -> ProjectionTest | None:
    return await Test.find_one().project(projection_model=ProjectionTest)


def find_one_update() -> UpdateOne:
    return Test.find_one().update({"$set": {"bar": "updated"}})


def find_one_upsert() -> UpdateOne:
    return Test.find_one().upsert(
        {"$set": {"bar": "updated"}},
        on_insert=Test(foo="foo", bar="bar", baz="baz"),
    )
