Refs #18586 -- Split custom_pk test.

1 files changed, 106 insertions, 55 deletions

diff --git a/tests/custom_pk/tests.py b/tests/custom_pk/tests.py
index 09a3ee5888..e3133774e6 100644
--- a/tests/custom_pk/tests.py
+++ b/tests/custom_pk/tests.py
@@ -8,40 +8,36 @@ from django.utils import six
 from .models import Employee, Business, Bar, Foo
 
 
-class CustomPKTests(TestCase):
-    def test_custom_pk(self):
-        dan = Employee.objects.create(
-            employee_code=123, first_name="Dan", last_name="Jones"
-        )
+class BasicCustomPKTests(TestCase):
+    @classmethod
+    def setUpTestData(cls):
+        cls.dan = Employee.objects.create(
+            employee_code=123, first_name="Dan", last_name="Jones",
+        )
+        cls.fran = Employee.objects.create(
+            employee_code=456, first_name="Fran", last_name="Bones",
+        )
+        cls.business = Business.objects.create(name="Sears")
+        cls.business.employees.add(cls.dan, cls.fran)
+
+    def test_querysets(self):
+        """
+        Both pk and custom attribute_name can be used in filter and friends
+        """
         self.assertQuerysetEqual(
-            Employee.objects.all(), [
+            Employee.objects.filter(pk=123), [
                 "Dan Jones",
             ],
             six.text_type
         )
 
-        fran = Employee.objects.create(
-            employee_code=456, first_name="Fran", last_name="Bones"
-        )
         self.assertQuerysetEqual(
-            Employee.objects.all(), [
-                "Fran Bones",
+            Employee.objects.filter(employee_code=123), [
                 "Dan Jones",
             ],
             six.text_type
         )
 
-        self.assertEqual(Employee.objects.get(pk=123), dan)
-        self.assertEqual(Employee.objects.get(pk=456), fran)
-
-        self.assertRaises(
-            Employee.DoesNotExist,
-            lambda: Employee.objects.get(pk=42)
-        )
-
-        # Use the name of the primary key, rather than pk.
-        self.assertEqual(Employee.objects.get(employee_code=123), dan)
-        # pk can be used as a substitute for the primary key.
         self.assertQuerysetEqual(
             Employee.objects.filter(pk__in=[123, 456]), [
                 "Fran Bones",
@@ -49,73 +45,61 @@ class CustomPKTests(TestCase):
             ],
             six.text_type
         )
-        # The primary key can be accessed via the pk property on the model.
-        e = Employee.objects.get(pk=123)
-        self.assertEqual(e.pk, 123)
-        # Or we can use the real attribute name for the primary key:
-        self.assertEqual(e.employee_code, 123)
-
-        # Fran got married and changed her last name.
-        fran = Employee.objects.get(pk=456)
-        fran.last_name = "Jones"
-        fran.save()
 
         self.assertQuerysetEqual(
-            Employee.objects.filter(last_name="Jones"), [
+            Employee.objects.all(), [
+                "Fran Bones",
                 "Dan Jones",
-                "Fran Jones",
             ],
             six.text_type
         )
 
-        emps = Employee.objects.in_bulk([123, 456])
-        self.assertEqual(emps[123], dan)
-
-        b = Business.objects.create(name="Sears")
-        b.employees.add(dan, fran)
         self.assertQuerysetEqual(
-            b.employees.all(), [
-                "Dan Jones",
-                "Fran Jones",
+            Business.objects.filter(name="Sears"), [
+                "Sears"
             ],
-            six.text_type
+            lambda b: b.name
         )
         self.assertQuerysetEqual(
-            fran.business_set.all(), [
+            Business.objects.filter(pk="Sears"), [
                 "Sears",
             ],
             lambda b: b.name
         )
 
-        self.assertEqual(Business.objects.in_bulk(["Sears"]), {
-            "Sears": b,
-        })
-
+    def test_querysets_related_name(self):
+        """
+        Custom pk doesn't affect related_name based lookups
+        """
         self.assertQuerysetEqual(
-            Business.objects.filter(name="Sears"), [
-                "Sears"
+            self.business.employees.all(), [
+                "Fran Bones",
+                "Dan Jones",
             ],
-            lambda b: b.name
+            six.text_type
         )
         self.assertQuerysetEqual(
-            Business.objects.filter(pk="Sears"), [
+            self.fran.business_set.all(), [
                 "Sears",
             ],
             lambda b: b.name
         )
 
-        # Queries across tables, involving primary key
+    def test_querysets_relational(self):
+        """
+        Queries across tables, involving primary key
+        """
         self.assertQuerysetEqual(
             Employee.objects.filter(business__name="Sears"), [
+                "Fran Bones",
                 "Dan Jones",
-                "Fran Jones",
             ],
             six.text_type,
         )
         self.assertQuerysetEqual(
             Employee.objects.filter(business__pk="Sears"), [
+                "Fran Bones",
                 "Dan Jones",
-                "Fran Jones",
             ],
             six.text_type,
         )
@@ -140,6 +124,73 @@ class CustomPKTests(TestCase):
             lambda b: b.name
         )
 
+    def test_get(self):
+        """
+        Get can accept pk or the real attribute name
+        """
+        self.assertEqual(Employee.objects.get(pk=123), self.dan)
+        self.assertEqual(Employee.objects.get(pk=456), self.fran)
+
+        self.assertRaises(
+            Employee.DoesNotExist,
+            lambda: Employee.objects.get(pk=42)
+        )
+
+        # Use the name of the primary key, rather than pk.
+        self.assertEqual(Employee.objects.get(employee_code=123), self.dan)
+
+    def test_pk_attributes(self):
+        """
+        pk and attribute name are available on the model
+        No default id attribute is added
+        """
+        # pk can be used as a substitute for the primary key.
+        # The primary key can be accessed via the pk property on the model.
+        e = Employee.objects.get(pk=123)
+        self.assertEqual(e.pk, 123)
+        # Or we can use the real attribute name for the primary key:
+        self.assertEqual(e.employee_code, 123)
+
+        self.assertRaises(AttributeError, lambda: e.id)
+
+    def test_in_bulk(self):
+        """
+        Custom pks work with in_bulk, both for integer and non-integer types
+        """
+        emps = Employee.objects.in_bulk([123, 456])
+        self.assertEqual(emps[123], self.dan)
+
+        self.assertEqual(Business.objects.in_bulk(["Sears"]), {
+            "Sears": self.business,
+        })
+
+    def test_save(self):
+        """
+        custom pks do not affect save
+        """
+        fran = Employee.objects.get(pk=456)
+        fran.last_name = "Jones"
+        fran.save()
+
+        self.assertQuerysetEqual(
+            Employee.objects.filter(last_name="Jones"), [
+                "Dan Jones",
+                "Fran Jones",
+            ],
+            six.text_type
+        )
+
+
+class CustomPKTests(TestCase):
+    def test_custom_pk_create(self):
+        """
+        New objects can be created both with pk and the custom name
+        """
+        Employee.objects.create(employee_code=1234, first_name="Foo", last_name="Bar")
+        Employee.objects.create(pk=1235, first_name="Foo", last_name="Baz")
+        Business.objects.create(name="Bears")
+        Business.objects.create(pk="Tears")
+
     def test_unicode_pk(self):
         # Primary key may be unicode string
         Business.objects.create(name='jaźń')