from django.test import TestCase

from .models import Flea, House, Person, Pet, Room

class UUIDPrefetchRelated(TestCase):

    def test_prefetch_related_from_uuid_model(self):

        Pet.objects.create(name="Fifi").people.add(

            Person.objects.create(name="Ellen"),

            Person.objects.create(name="George"),

        )

        with self.assertNumQueries(2):

            pet = Pet.objects.prefetch_related("people").get(name="Fifi")

        with self.assertNumQueries(0):

            self.assertEqual(2, len(pet.people.all()))

    def test_prefetch_related_to_uuid_model(self):

        Person.objects.create(name="Bella").pets.add(

            Pet.objects.create(name="Socks"),

            Pet.objects.create(name="Coffee"),

        )

        with self.assertNumQueries(2):

            person = Person.objects.prefetch_related("pets").get(name="Bella")

        with self.assertNumQueries(0):

            self.assertEqual(2, len(person.pets.all()))

    def test_prefetch_related_from_uuid_model_to_uuid_model(self):

        fleas = [Flea.objects.create() for i in range(3)]

        Pet.objects.create(name="Fifi").fleas_hosted.add(*fleas)

        Pet.objects.create(name="Bobo").fleas_hosted.add(*fleas)

        with self.assertNumQueries(2):

            pet = Pet.objects.prefetch_related("fleas_hosted").get(name="Fifi")

        with self.assertNumQueries(0):

            self.assertEqual(3, len(pet.fleas_hosted.all()))

        with self.assertNumQueries(2):

            flea = Flea.objects.prefetch_related("pets_visited").get(pk=fleas[0].pk)

        with self.assertNumQueries(0):

            self.assertEqual(2, len(flea.pets_visited.all()))

    def test_prefetch_related_from_uuid_model_to_uuid_model_with_values_flat(self):

        pet = Pet.objects.create(name="Fifi")

        pet.people.add(

            Person.objects.create(name="Ellen"),

            Person.objects.create(name="George"),

        )

        self.assertSequenceEqual(

            Pet.objects.prefetch_related("fleas_hosted").values_list("id", flat=True),

            [pet.id],

        )

class UUIDPrefetchRelatedLookups(TestCase):

    @classmethod

    def setUpTestData(cls):

        house = House.objects.create(name="Redwood", address="Arcata")

        room = Room.objects.create(name="Racoon", house=house)

        fleas = [Flea.objects.create(current_room=room) for i in range(3)]

        pet = Pet.objects.create(name="Spooky")

        pet.fleas_hosted.add(*fleas)

        person = Person.objects.create(name="Bob")

        person.houses.add(house)

        person.pets.add(pet)

        person.fleas_hosted.add(*fleas)

    def test_from_uuid_pk_lookup_uuid_pk_integer_pk(self):

        # From uuid-pk model, prefetch <uuid-pk model>.<integer-pk model>:

        with self.assertNumQueries(4):

            spooky = Pet.objects.prefetch_related(

                "fleas_hosted__current_room__house"

            ).get(name="Spooky")

        with self.assertNumQueries(0):

            self.assertEqual("Racoon", spooky.fleas_hosted.all()[0].current_room.name)

    def test_from_uuid_pk_lookup_integer_pk2_uuid_pk2(self):

        # From uuid-pk model, prefetch

        # <integer-pk model>.<integer-pk model>.<uuid-pk model>.<uuid-pk model>:

        with self.assertNumQueries(5):

            spooky = Pet.objects.prefetch_related("people__houses__rooms__fleas").get(

                name="Spooky"

            )

        with self.assertNumQueries(0):

            self.assertEqual(

                3,

                len(spooky.people.all()[0].houses.all()[0].rooms.all()[0].fleas.all()),

            )

    def test_from_integer_pk_lookup_uuid_pk_integer_pk(self):

        # From integer-pk model, prefetch <uuid-pk model>.<integer-pk model>:

        with self.assertNumQueries(3):

            racoon = Room.objects.prefetch_related("fleas__people_visited").get(

                name="Racoon"

            )

        with self.assertNumQueries(0):

            self.assertEqual("Bob", racoon.fleas.all()[0].people_visited.all()[0].name)

    def test_from_integer_pk_lookup_integer_pk_uuid_pk(self):

        # From integer-pk model, prefetch <integer-pk model>.<uuid-pk model>:

        with self.assertNumQueries(3):

            redwood = House.objects.prefetch_related("rooms__fleas").get(name="Redwood")

        with self.assertNumQueries(0):

            self.assertEqual(3, len(redwood.rooms.all()[0].fleas.all()))

    def test_from_integer_pk_lookup_integer_pk_uuid_pk_uuid_pk(self):

        # From integer-pk model, prefetch

        # <integer-pk model>.<uuid-pk model>.<uuid-pk model>:

        with self.assertNumQueries(4):

            redwood = House.objects.prefetch_related("rooms__fleas__pets_visited").get(

                name="Redwood"

            )

        with self.assertNumQueries(0):

            self.assertEqual(

                "Spooky",

                redwood.rooms.all()[0].fleas.all()[0].pets_visited.all()[0].name,

            )