from django.db.transaction import atomic
from django.test import TransactionTestCase

from .models import City, Country, Person, PersonProfile


class ForeignKeySelectForUpdateTestCase(TransactionTestCase):
    def setUp(self):
        Country.objects.create(pk=1, name="Belgium")
        Country.objects.create(pk=2, name="France")
        City.objects.bulk_create(
            (
                City(pk=1, name="Liberchies", country_id=1),
                City(pk=2, name="Samois-sur-Seine", country_id=2),
            )
        )
        Person.objects.bulk_create(
            (
                Person(pk=3, name="Reinhardt", born_id=1, died_id=2),
                Person(pk=4, name="Lon Sonterre", born_id=2),
            )
        )
        PersonProfile.objects.bulk_create(
            (PersonProfile(pk=3, person_id=4), PersonProfile(pk=5, person_id=3))
        )

    def test_select_for_update_reverse_one_to_one(self):
        with atomic():
            data = [
                *Person.objects.values_list("profile__pk", flat=True)
                .order_by("profile__pk")
                .select_for_update(of=("self",))
                .filter(pk=3)
            ]
        self.assertListEqual(data, [5])

    def test_select_for_update_forward_nullable_foreign_key(self):
        with atomic():
            data = [
                *Person.objects.select_for_update(of=("self",))
                .values_list("died__name", flat=True)
                .order_by("died__name")
                .filter(pk=3)
            ]
        self.assertListEqual(data, ["Samois-sur-Seine"])