Преглед изворни кода

feat(models): Handle GFK attributes in CloningMixin

Extend the CloningMixin to inject GenericForeignKey (GFK) attributes
when both content type and ID fields are present. Improves support for
models using GFK fields during cloning operations.

Fixes #21201
Martin Hauser пре 2 недеља
родитељ
комит
b26c7f34cd
2 измењених фајлова са 69 додато и 2 уклоњено
  1. 8 1
      netbox/netbox/models/features.py
  2. 61 1
      netbox/netbox/tests/test_model_features.py

+ 8 - 1
netbox/netbox/models/features.py

@@ -2,7 +2,7 @@ import json
 from collections import defaultdict
 from collections import defaultdict
 from functools import cached_property
 from functools import cached_property
 
 
-from django.contrib.contenttypes.fields import GenericRelation
+from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelation
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.contenttypes.models import ContentType
 from django.core.validators import ValidationError
 from django.core.validators import ValidationError
 from django.db import models
 from django.db import models
@@ -161,6 +161,13 @@ class CloningMixin(models.Model):
             elif field_value not in (None, ''):
             elif field_value not in (None, ''):
                 attrs[field_name] = field_value
                 attrs[field_name] = field_value
 
 
+        # Handle GenericForeignKeys. If the CT and ID fields are being cloned, also
+        # include the name of the GFK attribute itself, as this is what forms expect.
+        for field in self._meta.private_fields:
+            if isinstance(field, GenericForeignKey):
+                if field.ct_field in attrs and field.fk_field in attrs:
+                    attrs[field.name] = attrs[field.fk_field]
+
         # Include tags (if applicable)
         # Include tags (if applicable)
         if is_taggable(self):
         if is_taggable(self):
             attrs['tags'] = [tag.pk for tag in self.tags.all()]
             attrs['tags'] = [tag.pk for tag in self.tags.all()]

+ 61 - 1
netbox/netbox/tests/test_model_features.py

@@ -1,18 +1,28 @@
+from unittest import skipIf
+
+from django.conf import settings
 from django.test import TestCase
 from django.test import TestCase
 
 
 from core.models import AutoSyncRecord, DataSource
 from core.models import AutoSyncRecord, DataSource
+from dcim.models import Site
 from extras.models import CustomLink
 from extras.models import CustomLink
+from ipam.models import Prefix
 from netbox.models.features import get_model_features, has_feature, model_is_public
 from netbox.models.features import get_model_features, has_feature, model_is_public
-from netbox.tests.dummy_plugin.models import DummyModel
 from taggit.models import Tag
 from taggit.models import Tag
 
 
 
 
 class ModelFeaturesTestCase(TestCase):
 class ModelFeaturesTestCase(TestCase):
+    """
+    A test case class for verifying model features and utility functions.
+    """
 
 
+    @skipIf('netbox.tests.dummy_plugin' not in settings.PLUGINS, 'dummy_plugin not in settings.PLUGINS')
     def test_model_is_public(self):
     def test_model_is_public(self):
         """
         """
         Test that the is_public() utility function returns True for public models only.
         Test that the is_public() utility function returns True for public models only.
         """
         """
+        from netbox.tests.dummy_plugin.models import DummyModel
+
         # Public model
         # Public model
         self.assertFalse(hasattr(DataSource, '_netbox_private'))
         self.assertFalse(hasattr(DataSource, '_netbox_private'))
         self.assertTrue(model_is_public(DataSource))
         self.assertTrue(model_is_public(DataSource))
@@ -51,3 +61,53 @@ class ModelFeaturesTestCase(TestCase):
         features = get_model_features(CustomLink)
         features = get_model_features(CustomLink)
         self.assertIn('cloning', features)
         self.assertIn('cloning', features)
         self.assertNotIn('bookmarks', features)
         self.assertNotIn('bookmarks', features)
+
+    def test_cloningmixin_injects_gfk_attribute(self):
+        """
+        Tests the cloning mixin with GFK attribute injection in the `clone` method.
+
+        This test validates that the `clone` method correctly handles
+        and retains the General Foreign Key (GFK) attributes on an
+        object when the cloning fields are explicitly defined.
+        """
+        site = Site.objects.create(name='Test Site', slug='test-site')
+        prefix = Prefix.objects.create(prefix='10.0.0.0/24', scope=site)
+
+        original_clone_fields = getattr(Prefix, 'clone_fields', None)
+        try:
+            Prefix.clone_fields = ('scope_type', 'scope_id')
+            attrs = prefix.clone()
+
+            self.assertEqual(attrs['scope_type'], prefix.scope_type_id)
+            self.assertEqual(attrs['scope_id'], prefix.scope_id)
+            self.assertEqual(attrs['scope'], prefix.scope_id)
+        finally:
+            if original_clone_fields is None:
+                delattr(Prefix, 'clone_fields')
+            else:
+                Prefix.clone_fields = original_clone_fields
+
+    def test_cloningmixin_does_not_inject_gfk_attribute_if_incomplete(self):
+        """
+        Tests the cloning mixin with incomplete cloning fields does not inject the GFK attribute.
+
+        This test validates that the `clone` method correctly handles
+        the case where the cloning fields are incomplete, ensuring that
+        the generic foreign key (GFK) attribute is not injected during
+        the cloning process.
+        """
+        site = Site.objects.create(name='Test Site', slug='test-site')
+        prefix = Prefix.objects.create(prefix='10.0.0.0/24', scope=site)
+
+        original_clone_fields = getattr(Prefix, 'clone_fields', None)
+        try:
+            Prefix.clone_fields = ('scope_type',)
+            attrs = prefix.clone()
+
+            self.assertIn('scope_type', attrs)
+            self.assertNotIn('scope', attrs)
+        finally:
+            if original_clone_fields is None:
+                delattr(Prefix, 'clone_fields')
+            else:
+                Prefix.clone_fields = original_clone_fields