169 lines
5.4 KiB
Python
169 lines
5.4 KiB
Python
# type: ignore
|
|
from typing import List, Optional
|
|
|
|
import databases
|
|
import pytest
|
|
import sqlalchemy as sa
|
|
from pydantic.typing import ForwardRef
|
|
from sqlalchemy import create_engine
|
|
|
|
import ormar
|
|
from ormar import ModelMeta
|
|
from tests.settings import DATABASE_URL
|
|
|
|
metadata = sa.MetaData()
|
|
db = databases.Database(DATABASE_URL)
|
|
engine = create_engine(DATABASE_URL)
|
|
|
|
TeacherRef = ForwardRef("Teacher")
|
|
|
|
|
|
class BaseMeta(ormar.ModelMeta):
|
|
metadata = metadata
|
|
database = db
|
|
|
|
|
|
class Student(ormar.Model):
|
|
class Meta(BaseMeta):
|
|
pass
|
|
|
|
id: int = ormar.Integer(primary_key=True)
|
|
name: str = ormar.String(max_length=100)
|
|
primary_teacher: TeacherRef = ormar.ForeignKey(
|
|
TeacherRef, related_name="own_students"
|
|
)
|
|
|
|
|
|
class StudentTeacher(ormar.Model):
|
|
class Meta(BaseMeta):
|
|
tablename = "students_x_teachers"
|
|
|
|
|
|
class Teacher(ormar.Model):
|
|
class Meta(ModelMeta):
|
|
metadata = metadata
|
|
database = db
|
|
|
|
id: int = ormar.Integer(primary_key=True)
|
|
name: str = ormar.String(max_length=100)
|
|
students = ormar.ManyToMany(
|
|
Student, through=StudentTeacher, related_name="teachers"
|
|
)
|
|
|
|
|
|
Student.update_forward_refs()
|
|
|
|
CityRef = ForwardRef("City")
|
|
CountryRef = ForwardRef("Country")
|
|
|
|
|
|
class Country(ormar.Model):
|
|
class Meta(BaseMeta):
|
|
tablename = "countries"
|
|
|
|
id: int = ormar.Integer(primary_key=True)
|
|
name: str = ormar.String(max_length=128)
|
|
capital: Optional[CityRef] = ormar.ForeignKey(
|
|
CityRef, related_name="capital_city", nullable=True
|
|
)
|
|
borders: Optional[List[CountryRef]] = ormar.ManyToMany(CountryRef)
|
|
|
|
|
|
class City(ormar.Model):
|
|
class Meta(BaseMeta):
|
|
tablename = "cities"
|
|
|
|
id: int = ormar.Integer(primary_key=True)
|
|
name: str = ormar.String(max_length=128)
|
|
country: Country = ormar.ForeignKey(
|
|
Country, related_name="cities", skip_reverse=True
|
|
)
|
|
|
|
|
|
Country.update_forward_refs()
|
|
|
|
|
|
@pytest.fixture(autouse=True, scope="module")
|
|
def create_test_database():
|
|
metadata.create_all(engine)
|
|
yield
|
|
metadata.drop_all(engine)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_double_relations():
|
|
async with db:
|
|
async with db.transaction(force_rollback=True):
|
|
t1 = await Teacher.objects.create(name="Mr. Jones")
|
|
t2 = await Teacher.objects.create(name="Ms. Smith")
|
|
t3 = await Teacher.objects.create(name="Mr. Quibble")
|
|
|
|
s1 = await Student.objects.create(name="Joe", primary_teacher=t1)
|
|
s2 = await Student.objects.create(name="Sam", primary_teacher=t1)
|
|
s3 = await Student.objects.create(name="Kate", primary_teacher=t2)
|
|
s4 = await Student.objects.create(name="Zoe", primary_teacher=t2)
|
|
s5 = await Student.objects.create(name="John", primary_teacher=t3)
|
|
s6 = await Student.objects.create(name="Anna", primary_teacher=t3)
|
|
|
|
for t in [t1, t2, t3]:
|
|
for s in [s1, s2, s3, s4, s5, s6]:
|
|
await t.students.add(s)
|
|
|
|
jones = (
|
|
await Teacher.objects.select_related(["students", "own_students"])
|
|
.order_by(["students__name", "own_students__name"])
|
|
.get(name="Mr. Jones")
|
|
)
|
|
assert len(jones.students) == 6
|
|
assert jones.students[0].name == "Anna"
|
|
assert jones.students[5].name == "Zoe"
|
|
assert len(jones.own_students) == 2
|
|
assert jones.own_students[0].name == "Joe"
|
|
assert jones.own_students[1].name == "Sam"
|
|
|
|
smith = (
|
|
await Teacher.objects.select_related(["students", "own_students"])
|
|
.filter(students__name__contains="a")
|
|
.order_by(["students__name", "own_students__name"])
|
|
.get(name="Ms. Smith")
|
|
)
|
|
assert len(smith.students) == 3
|
|
assert smith.students[0].name == "Anna"
|
|
assert smith.students[2].name == "Sam"
|
|
assert len(smith.own_students) == 2
|
|
assert smith.own_students[0].name == "Kate"
|
|
assert smith.own_students[1].name == "Zoe"
|
|
|
|
quibble = (
|
|
await Teacher.objects.select_related(["students", "own_students"])
|
|
.filter(students__name__startswith="J")
|
|
.order_by(["-students__name", "own_students__name"])
|
|
.get(name="Mr. Quibble")
|
|
)
|
|
assert len(quibble.students) == 2
|
|
assert quibble.students[1].name == "Joe"
|
|
assert quibble.students[0].name == "John"
|
|
assert len(quibble.own_students) == 2
|
|
assert quibble.own_students[1].name == "John"
|
|
assert quibble.own_students[0].name == "Anna"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_auto_through_model():
|
|
async with db:
|
|
async with db.transaction(force_rollback=True):
|
|
england = await Country(name="England").save()
|
|
france = await Country(name="France").save()
|
|
london = await City(name="London", country=england).save()
|
|
england.capital = london
|
|
await england.update()
|
|
await england.borders.add(france)
|
|
|
|
check = await Country.objects.select_related(["capital", "borders"]).get(
|
|
name="England"
|
|
)
|
|
assert check.name == "England"
|
|
assert check.capital.name == "London"
|
|
assert check.capital.country.pk == check.pk
|
|
assert check.borders[0] == france
|