Merge pull request #5192 from matteius/DRF-5135-one-to-one-pk

tomchristie · web-flow · commit 598e5877cd0b · 2017-06-16T12:06:12.000+01:00
Special case for when OneToOneField is also primary key.
diff --git a/rest_framework/serializers.py b/rest_framework/serializers.py
@@ -1159,6 +1159,11 @@ def build_standard_field(self, field_name, model_field):
         field_class = field_mapping[model_field]
         field_kwargs = get_field_kwargs(field_name, model_field)
 
+        # Special case to handle when a OneToOneField is also the primary key
+        if model_field.one_to_one and model_field.primary_key:
+            field_class = self.serializer_related_field
+            field_kwargs['queryset'] = model_field.related_model.objects
+
         if 'choices' in field_kwargs:
             # Fields with choices get coerced into `ChoiceField`
             # instead of using their regular typed field.
diff --git a/tests/models.py b/tests/models.py
@@ -88,3 +88,11 @@ class NullableOneToOneSource(RESTFrameworkModel):
     target = models.OneToOneField(
         OneToOneTarget, null=True, blank=True,
         related_name='nullable_source', on_delete=models.CASCADE)
+
+
+class OneToOnePKSource(RESTFrameworkModel):
+    """ Test model where the primary key is a OneToOneField with another model. """
+    name = models.CharField(max_length=100)
+    target = models.OneToOneField(
+        OneToOneTarget, primary_key=True,
+        related_name='required_source', on_delete=models.CASCADE)
diff --git a/tests/test_relations_pk.py b/tests/test_relations_pk.py
@@ -6,8 +6,8 @@
 from rest_framework import serializers
 from tests.models import (
     ForeignKeySource, ForeignKeyTarget, ManyToManySource, ManyToManyTarget,
-    NullableForeignKeySource, NullableOneToOneSource,
-    NullableUUIDForeignKeySource, OneToOneTarget, UUIDForeignKeyTarget
+    NullableForeignKeySource, NullableOneToOneSource, NullableUUIDForeignKeySource,
+    OneToOnePKSource, OneToOneTarget, UUIDForeignKeyTarget
 )
 
 
@@ -63,6 +63,13 @@ class Meta:
         fields = ('id', 'name', 'nullable_source')
 
 
+class OneToOnePKSourceSerializer(serializers.ModelSerializer):
+
+    class Meta:
+        model = OneToOnePKSource
+        fields = '__all__'
+
+
 # TODO: Add test that .data cannot be accessed prior to .is_valid
 
 class PKManyToManyTests(TestCase):
@@ -486,3 +493,51 @@ def test_reverse_foreign_key_retrieve_with_null(self):
             {'id': 2, 'name': 'target-2', 'nullable_source': 1},
         ]
         assert serializer.data == expected
+
+
+class OneToOnePrimaryKeyTests(TestCase):
+
+    def setUp(self):
+        # Given: Some target models already exist
+        self.target = target = OneToOneTarget(name='target-1')
+        target.save()
+        self.alt_target = alt_target = OneToOneTarget(name='target-2')
+        alt_target.save()
+
+    def test_one_to_one_when_primary_key(self):
+        # When: Creating a Source pointing at the id of the second Target
+        target_pk = self.alt_target.id
+        source = OneToOnePKSourceSerializer(data={'name': 'source-2', 'target': target_pk})
+        # Then: The source is valid with the serializer
+        if not source.is_valid():
+            self.fail("Expected OneToOnePKTargetSerializer to be valid but had errors: {}".format(source.errors))
+        # Then: Saving the serializer creates a new object
+        new_source = source.save()
+        # Then: The new object has the same pk as the target object
+        self.assertEqual(new_source.pk, target_pk)
+
+    def test_one_to_one_when_primary_key_no_duplicates(self):
+        # When: Creating a Source pointing at the id of the second Target
+        target_pk = self.target.id
+        data = {'name': 'source-1', 'target': target_pk}
+        source = OneToOnePKSourceSerializer(data=data)
+        # Then: The source is valid with the serializer
+        self.assertTrue(source.is_valid())
+        # Then: Saving the serializer creates a new object
+        new_source = source.save()
+        # Then: The new object has the same pk as the target object
+        self.assertEqual(new_source.pk, target_pk)
+        # When: Trying to create a second object
+        second_source = OneToOnePKSourceSerializer(data=data)
+        self.assertFalse(second_source.is_valid())
+        expected = {'target': [u'one to one pk source with this target already exists.']}
+        self.assertDictEqual(second_source.errors, expected)
+
+    def test_one_to_one_when_primary_key_does_not_exist(self):
+        # Given: a target PK that does not exist
+        target_pk = self.target.pk + self.alt_target.pk
+        source = OneToOnePKSourceSerializer(data={'name': 'source-2', 'target': target_pk})
+        # Then: The source is not valid with the serializer
+        self.assertFalse(source.is_valid())
+        self.assertIn("Invalid pk", source.errors['target'][0])
+        self.assertIn("object does not exist", source.errors['target'][0])
