from typing import List, Optional import ormar import pytest from tests.lifespan import init_tests from tests.settings import create_config base_ormar_config = create_config() class Child(ormar.Model): ormar_config = base_ormar_config.copy(tablename="children") id: int = ormar.Integer(name="child_id", primary_key=True) first_name: str = ormar.String(name="fname", max_length=100) last_name: str = ormar.String(name="lname", max_length=100) born_year: int = ormar.Integer(name="year_born", nullable=True) class Artist(ormar.Model): ormar_config = base_ormar_config.copy(tablename="artists") id: int = ormar.Integer(name="artist_id", primary_key=True) first_name: str = ormar.String(name="fname", max_length=100) last_name: str = ormar.String(name="lname", max_length=100) born_year: int = ormar.Integer(name="year") children: Optional[List[Child]] = ormar.ManyToMany(Child) class Album(ormar.Model): ormar_config = base_ormar_config.copy(tablename="music_albums") id: int = ormar.Integer(name="album_id", primary_key=True) name: str = ormar.String(name="album_name", max_length=100) artist: Optional[Artist] = ormar.ForeignKey(Artist, name="artist_id") create_test_database = init_tests(base_ormar_config) def test_table_structure(): assert "album_id" in [x.name for x in Album.ormar_config.table.columns] assert "album_name" in [x.name for x in Album.ormar_config.table.columns] assert "fname" in [x.name for x in Artist.ormar_config.table.columns] assert "lname" in [x.name for x in Artist.ormar_config.table.columns] assert "year" in [x.name for x in Artist.ormar_config.table.columns] @pytest.mark.asyncio async def test_working_with_aliases(): async with base_ormar_config.database: async with base_ormar_config.database.transaction(force_rollback=True): artist = await Artist.objects.create( first_name="Ted", last_name="Mosbey", born_year=1975 ) await Album.objects.create(name="Aunt Robin", artist=artist) await artist.children.create( first_name="Son", last_name="1", born_year=1990 ) await artist.children.create( first_name="Son", last_name="2", born_year=1995 ) await artist.children.create( first_name="Son", last_name="3", born_year=1998 ) album = await Album.objects.select_related("artist").first() assert album.artist.last_name == "Mosbey" assert album.artist.id is not None assert album.artist.first_name == "Ted" assert album.artist.born_year == 1975 assert album.name == "Aunt Robin" artist = await Artist.objects.select_related("children").get() assert len(artist.children) == 3 assert artist.children[0].first_name == "Son" assert artist.children[1].last_name == "2" assert artist.children[2].last_name == "3" await artist.update(last_name="Bundy") await Artist.objects.filter(pk=artist.pk).update(born_year=1974) artist = await Artist.objects.select_related("children").get() assert artist.last_name == "Bundy" assert artist.born_year == 1974 artist = ( await Artist.objects.select_related("children") .fields( [ "first_name", "last_name", "born_year", "children__first_name", "children__last_name", ] ) .get() ) assert artist.children[0].born_year is None @pytest.mark.asyncio async def test_bulk_operations_and_fields(): async with base_ormar_config.database: d1 = Child(first_name="Daughter", last_name="1", born_year=1990) d2 = Child(first_name="Daughter", last_name="2", born_year=1991) await Child.objects.bulk_create([d1, d2]) children = await Child.objects.filter(first_name="Daughter").all() assert len(children) == 2 assert children[0].last_name == "1" for child in children: child.born_year = child.born_year - 100 await Child.objects.bulk_update(children) children = await Child.objects.filter(first_name="Daughter").all() assert len(children) == 2 assert children[0].born_year == 1890 children = await Child.objects.fields(["first_name", "last_name"]).all() assert len(children) == 2 for child in children: assert child.born_year is None await children[0].load() await children[0].delete() children = await Child.objects.all() assert len(children) == 1 @pytest.mark.asyncio async def test_working_with_aliases_get_or_create(): async with base_ormar_config.database: async with base_ormar_config.database.transaction(force_rollback=True): artist, created = await Artist.objects.get_or_create( first_name="Teddy", last_name="Bear", born_year=2020 ) assert artist.pk is not None assert created is True artist2, created = await Artist.objects.get_or_create( first_name="Teddy", last_name="Bear", born_year=2020 ) assert artist == artist2 assert created is False art3 = artist2.model_dump() art3["born_year"] = 2019 await Artist.objects.update_or_create(**art3) artist3 = await Artist.objects.get(last_name="Bear") assert artist3.born_year == 2019 artists = await Artist.objects.all() assert len(artists) == 1