1. from django.db import IntegrityError, transaction
    
  2. from django.test import TestCase, skipIfDBFeature, skipUnlessDBFeature
    
  3. 
    
  4. from .fields import MyWrapper
    
  5. from .models import Bar, Business, CustomAutoFieldModel, Employee, Foo
    
  6. 
    
  7. 
    
  8. class BasicCustomPKTests(TestCase):
    
  9.     @classmethod
    
  10.     def setUpTestData(cls):
    
  11.         cls.dan = Employee.objects.create(
    
  12.             employee_code=123,
    
  13.             first_name="Dan",
    
  14.             last_name="Jones",
    
  15.         )
    
  16.         cls.fran = Employee.objects.create(
    
  17.             employee_code=456,
    
  18.             first_name="Fran",
    
  19.             last_name="Bones",
    
  20.         )
    
  21.         cls.business = Business.objects.create(name="Sears")
    
  22.         cls.business.employees.add(cls.dan, cls.fran)
    
  23. 
    
  24.     def test_querysets(self):
    
  25.         """
    
  26.         Both pk and custom attribute_name can be used in filter and friends
    
  27.         """
    
  28.         self.assertQuerysetEqual(
    
  29.             Employee.objects.filter(pk=123),
    
  30.             [
    
  31.                 "Dan Jones",
    
  32.             ],
    
  33.             str,
    
  34.         )
    
  35. 
    
  36.         self.assertQuerysetEqual(
    
  37.             Employee.objects.filter(employee_code=123),
    
  38.             [
    
  39.                 "Dan Jones",
    
  40.             ],
    
  41.             str,
    
  42.         )
    
  43. 
    
  44.         self.assertQuerysetEqual(
    
  45.             Employee.objects.filter(pk__in=[123, 456]),
    
  46.             [
    
  47.                 "Fran Bones",
    
  48.                 "Dan Jones",
    
  49.             ],
    
  50.             str,
    
  51.         )
    
  52. 
    
  53.         self.assertQuerysetEqual(
    
  54.             Employee.objects.all(),
    
  55.             [
    
  56.                 "Fran Bones",
    
  57.                 "Dan Jones",
    
  58.             ],
    
  59.             str,
    
  60.         )
    
  61. 
    
  62.         self.assertQuerysetEqual(
    
  63.             Business.objects.filter(name="Sears"), ["Sears"], lambda b: b.name
    
  64.         )
    
  65.         self.assertQuerysetEqual(
    
  66.             Business.objects.filter(pk="Sears"),
    
  67.             [
    
  68.                 "Sears",
    
  69.             ],
    
  70.             lambda b: b.name,
    
  71.         )
    
  72. 
    
  73.     def test_querysets_related_name(self):
    
  74.         """
    
  75.         Custom pk doesn't affect related_name based lookups
    
  76.         """
    
  77.         self.assertQuerysetEqual(
    
  78.             self.business.employees.all(),
    
  79.             [
    
  80.                 "Fran Bones",
    
  81.                 "Dan Jones",
    
  82.             ],
    
  83.             str,
    
  84.         )
    
  85.         self.assertQuerysetEqual(
    
  86.             self.fran.business_set.all(),
    
  87.             [
    
  88.                 "Sears",
    
  89.             ],
    
  90.             lambda b: b.name,
    
  91.         )
    
  92. 
    
  93.     def test_querysets_relational(self):
    
  94.         """
    
  95.         Queries across tables, involving primary key
    
  96.         """
    
  97.         self.assertQuerysetEqual(
    
  98.             Employee.objects.filter(business__name="Sears"),
    
  99.             [
    
  100.                 "Fran Bones",
    
  101.                 "Dan Jones",
    
  102.             ],
    
  103.             str,
    
  104.         )
    
  105.         self.assertQuerysetEqual(
    
  106.             Employee.objects.filter(business__pk="Sears"),
    
  107.             [
    
  108.                 "Fran Bones",
    
  109.                 "Dan Jones",
    
  110.             ],
    
  111.             str,
    
  112.         )
    
  113. 
    
  114.         self.assertQuerysetEqual(
    
  115.             Business.objects.filter(employees__employee_code=123),
    
  116.             [
    
  117.                 "Sears",
    
  118.             ],
    
  119.             lambda b: b.name,
    
  120.         )
    
  121.         self.assertQuerysetEqual(
    
  122.             Business.objects.filter(employees__pk=123),
    
  123.             [
    
  124.                 "Sears",
    
  125.             ],
    
  126.             lambda b: b.name,
    
  127.         )
    
  128. 
    
  129.         self.assertQuerysetEqual(
    
  130.             Business.objects.filter(employees__first_name__startswith="Fran"),
    
  131.             [
    
  132.                 "Sears",
    
  133.             ],
    
  134.             lambda b: b.name,
    
  135.         )
    
  136. 
    
  137.     def test_get(self):
    
  138.         """
    
  139.         Get can accept pk or the real attribute name
    
  140.         """
    
  141.         self.assertEqual(Employee.objects.get(pk=123), self.dan)
    
  142.         self.assertEqual(Employee.objects.get(pk=456), self.fran)
    
  143. 
    
  144.         with self.assertRaises(Employee.DoesNotExist):
    
  145.             Employee.objects.get(pk=42)
    
  146. 
    
  147.         # Use the name of the primary key, rather than pk.
    
  148.         self.assertEqual(Employee.objects.get(employee_code=123), self.dan)
    
  149. 
    
  150.     def test_pk_attributes(self):
    
  151.         """
    
  152.         pk and attribute name are available on the model
    
  153.         No default id attribute is added
    
  154.         """
    
  155.         # pk can be used as a substitute for the primary key.
    
  156.         # The primary key can be accessed via the pk property on the model.
    
  157.         e = Employee.objects.get(pk=123)
    
  158.         self.assertEqual(e.pk, 123)
    
  159.         # Or we can use the real attribute name for the primary key:
    
  160.         self.assertEqual(e.employee_code, 123)
    
  161. 
    
  162.         with self.assertRaisesMessage(
    
  163.             AttributeError, "'Employee' object has no attribute 'id'"
    
  164.         ):
    
  165.             e.id
    
  166. 
    
  167.     def test_in_bulk(self):
    
  168.         """
    
  169.         Custom pks work with in_bulk, both for integer and non-integer types
    
  170.         """
    
  171.         emps = Employee.objects.in_bulk([123, 456])
    
  172.         self.assertEqual(emps[123], self.dan)
    
  173. 
    
  174.         self.assertEqual(
    
  175.             Business.objects.in_bulk(["Sears"]),
    
  176.             {
    
  177.                 "Sears": self.business,
    
  178.             },
    
  179.         )
    
  180. 
    
  181.     def test_save(self):
    
  182.         """
    
  183.         custom pks do not affect save
    
  184.         """
    
  185.         fran = Employee.objects.get(pk=456)
    
  186.         fran.last_name = "Jones"
    
  187.         fran.save()
    
  188. 
    
  189.         self.assertQuerysetEqual(
    
  190.             Employee.objects.filter(last_name="Jones"),
    
  191.             [
    
  192.                 "Dan Jones",
    
  193.                 "Fran Jones",
    
  194.             ],
    
  195.             str,
    
  196.         )
    
  197. 
    
  198. 
    
  199. class CustomPKTests(TestCase):
    
  200.     def test_custom_pk_create(self):
    
  201.         """
    
  202.         New objects can be created both with pk and the custom name
    
  203.         """
    
  204.         Employee.objects.create(employee_code=1234, first_name="Foo", last_name="Bar")
    
  205.         Employee.objects.create(pk=1235, first_name="Foo", last_name="Baz")
    
  206.         Business.objects.create(name="Bears")
    
  207.         Business.objects.create(pk="Tears")
    
  208. 
    
  209.     def test_unicode_pk(self):
    
  210.         # Primary key may be Unicode string.
    
  211.         Business.objects.create(name="jaźń")
    
  212. 
    
  213.     def test_unique_pk(self):
    
  214.         # The primary key must also be unique, so trying to create a new object
    
  215.         # with the same primary key will fail.
    
  216.         Employee.objects.create(
    
  217.             employee_code=123, first_name="Frank", last_name="Jones"
    
  218.         )
    
  219.         with self.assertRaises(IntegrityError):
    
  220.             with transaction.atomic():
    
  221.                 Employee.objects.create(
    
  222.                     employee_code=123, first_name="Fred", last_name="Jones"
    
  223.                 )
    
  224. 
    
  225.     def test_zero_non_autoincrement_pk(self):
    
  226.         Employee.objects.create(employee_code=0, first_name="Frank", last_name="Jones")
    
  227.         employee = Employee.objects.get(pk=0)
    
  228.         self.assertEqual(employee.employee_code, 0)
    
  229. 
    
  230.     def test_custom_field_pk(self):
    
  231.         # Regression for #10785 -- Custom fields can be used for primary keys.
    
  232.         new_bar = Bar.objects.create()
    
  233.         new_foo = Foo.objects.create(bar=new_bar)
    
  234. 
    
  235.         f = Foo.objects.get(bar=new_bar.pk)
    
  236.         self.assertEqual(f, new_foo)
    
  237.         self.assertEqual(f.bar, new_bar)
    
  238. 
    
  239.         f = Foo.objects.get(bar=new_bar)
    
  240.         self.assertEqual(f, new_foo),
    
  241.         self.assertEqual(f.bar, new_bar)
    
  242. 
    
  243.     # SQLite lets objects be saved with an empty primary key, even though an
    
  244.     # integer is expected. So we can't check for an error being raised in that
    
  245.     # case for SQLite. Remove it from the suite for this next bit.
    
  246.     @skipIfDBFeature("supports_unspecified_pk")
    
  247.     def test_required_pk(self):
    
  248.         # The primary key must be specified, so an error is raised if you
    
  249.         # try to create an object without it.
    
  250.         with self.assertRaises(IntegrityError):
    
  251.             with transaction.atomic():
    
  252.                 Employee.objects.create(first_name="Tom", last_name="Smith")
    
  253. 
    
  254.     def test_auto_field_subclass_create(self):
    
  255.         obj = CustomAutoFieldModel.objects.create()
    
  256.         self.assertIsInstance(obj.id, MyWrapper)
    
  257. 
    
  258.     @skipUnlessDBFeature("can_return_rows_from_bulk_insert")
    
  259.     def test_auto_field_subclass_bulk_create(self):
    
  260.         obj = CustomAutoFieldModel()
    
  261.         CustomAutoFieldModel.objects.bulk_create([obj])
    
  262.         self.assertIsInstance(obj.id, MyWrapper)