from django.contrib.auth import get_user_model
from django.contrib.auth.models import Group
from django.db.transaction import atomic
from django.test import TransactionTestCase

User = get_user_model()


class ProofOfConceptTest(TransactionTestCase):
    def setUp(self):
        super().setUp()
        groups = Group.objects.bulk_create(
            (
                Group(pk=1, name="Group 1"),
                Group(pk=2, name="Group 2"),
                Group(pk=3, name="Group 3"),
            )
        )
        users = User.objects.bulk_create(
            (
                User(
                    pk=1,
                    password="!",
                    first_name="Camilla",
                    last_name="Millbern",
                    email="cmillbern@aol.com",
                    username="cmillbern",
                ),
                User(
                    pk=2,
                    password="!",
                    first_name="Shaina",
                    last_name="Sonnen",
                    email="shainasonnen@gmail.com",
                    username="ssonnen43",
                ),
            )
        )
        users[0].groups.set(groups)
        users[1].groups.set(groups[:2])

    def test_always_works(self):
        with atomic():
            groups = [
                *User.objects.values_list("pk", "groups__name")
                .order_by("groups__name")
                .select_for_update(of=("self",))
                .filter(pk=1)
            ]
        self.assertSequenceEqual(
            groups, ((1, "Group 1"), (1, "Group 2"), (1, "Group 3"))
        )

    def test_proof_of_concept(self):
        with atomic():
            groups = [
                *User.objects.values_list("groups__name", flat=True)
                .order_by("groups__name")
                .select_for_update(of=("self",))
                .filter(pk=1)
            ]
        self.assertSequenceEqual(groups, ("Group 1", "Group 2", "Group 3"))