Просмотр исходного кода

Employ signals to update child objects when RackGroup/Rack site assignment changes

Jeremy Stretch 5 лет назад
Родитель
Сommit
e4f22bc494
3 измененных файлов с 102 добавлено и 17 удалено
  1. 0 16
      netbox/dcim/models/racks.py
  2. 41 1
      netbox/dcim/signals.py
  3. 61 0
      netbox/dcim/tests/test_models.py

+ 0 - 16
netbox/dcim/models/racks.py

@@ -326,22 +326,6 @@ class Rack(ChangeLoggedModel, CustomFieldModel):
                         'group': "Rack group must be from the same site, {}.".format(self.site)
                     })
 
-    def save(self, *args, **kwargs):
-
-        # Record the original site assignment for this rack.
-        _site_id = None
-        if self.pk:
-            _site_id = Rack.objects.get(pk=self.pk).site_id
-
-        super().save(*args, **kwargs)
-
-        # Update racked devices if the assigned Site has been changed.
-        if _site_id is not None and self.site_id != _site_id:
-            devices = Device.objects.filter(rack=self)
-            for device in devices:
-                device.site = self.site
-                device.save()
-
     def to_csv(self):
         return (
             self.site.name,

+ 41 - 1
netbox/dcim/signals.py

@@ -2,12 +2,13 @@ import logging
 
 from cacheops import invalidate_obj
 from django.contrib.contenttypes.models import ContentType
+from django.db.models import Q
 from django.db.models.signals import post_save, post_delete, pre_delete
 from django.db import transaction
 from django.dispatch import receiver
 
 from .choices import CableStatusChoices
-from .models import Cable, CablePath, Device, PathEndpoint, VirtualChassis
+from .models import Cable, CablePath, Device, PathEndpoint, Rack, RackGroup, VirtualChassis
 
 
 def create_cablepath(node):
@@ -36,6 +37,40 @@ def rebuild_paths(obj):
             create_cablepath(cp.origin)
 
 
+#
+# Site/rack/device assignment
+#
+
+@receiver(post_save, sender=RackGroup)
+def handle_rackgroup_site_change(instance, created, **kwargs):
+    """
+    Update child RackGroups and Racks if Site assignment has changed. We intentionally recurse through each child
+    object instead of calling update() on the QuerySet to ensure the proper change records get created for each.
+    """
+    if not created:
+        for rackgroup in instance.get_children():
+            rackgroup.site = instance.site
+            rackgroup.save()
+        for rack in Rack.objects.filter(group=instance).exclude(site=instance.site):
+            rack.site = instance.site
+            rack.save()
+
+
+@receiver(post_save, sender=Rack)
+def handle_rack_site_change(instance, created, **kwargs):
+    """
+    Update child Devices if Site assignment has changed.
+    """
+    if not created:
+        for device in Device.objects.filter(rack=instance).exclude(site=instance.site):
+            device.site = instance.site
+            device.save()
+
+
+#
+# Virtual chassis
+#
+
 @receiver(post_save, sender=VirtualChassis)
 def assign_virtualchassis_master(instance, created, **kwargs):
     """
@@ -60,6 +95,11 @@ def clear_virtualchassis_members(instance, **kwargs):
         device.save()
 
 
+#
+# Cables
+#
+
+
 @receiver(post_save, sender=Cable)
 def update_connected_endpoints(instance, created, raw=False, **kwargs):
     """

+ 61 - 0
netbox/dcim/tests/test_models.py

@@ -7,6 +7,39 @@ from dcim.models import *
 from tenancy.models import Tenant
 
 
+class RackGroupTestCase(TestCase):
+
+    def test_change_rackgroup_site(self):
+        """
+        Check that all child RackGroups and Racks get updated when a RackGroup is moved to a new Site. Topology:
+        Site A
+          - RackGroup A1
+            - RackGroup A2
+              - Rack 2
+            - Rack 1
+        """
+        site_a = Site.objects.create(name='Site A', slug='site-a')
+        site_b = Site.objects.create(name='Site B', slug='site-b')
+
+        rackgroup_a1 = RackGroup(site=site_a, name='RackGroup A1', slug='rackgroup-a1')
+        rackgroup_a1.save()
+        rackgroup_a2 = RackGroup(site=site_a, parent=rackgroup_a1, name='RackGroup A2', slug='rackgroup-a2')
+        rackgroup_a2.save()
+
+        rack1 = Rack.objects.create(site=site_a, group=rackgroup_a1, name='Rack 1')
+        rack2 = Rack.objects.create(site=site_a, group=rackgroup_a2, name='Rack 2')
+
+        # Move RackGroup A1 to Site B
+        rackgroup_a1.site = site_b
+        rackgroup_a1.save()
+
+        # Check that all objects within RackGroup A1 now belong to Site B
+        self.assertEqual(RackGroup.objects.get(pk=rackgroup_a1.pk).site, site_b)
+        self.assertEqual(RackGroup.objects.get(pk=rackgroup_a2.pk).site, site_b)
+        self.assertEqual(Rack.objects.get(pk=rack1.pk).site, site_b)
+        self.assertEqual(Rack.objects.get(pk=rack2.pk).site, site_b)
+
+
 class RackTestCase(TestCase):
 
     def setUp(self):
@@ -154,6 +187,34 @@ class RackTestCase(TestCase):
         )
         self.assertTrue(pdu)
 
+    def test_change_rack_site(self):
+        """
+        Check that child Devices get updated when a Rack is moved to a new Site.
+        """
+        site_a = Site.objects.create(name='Site A', slug='site-a')
+        site_b = Site.objects.create(name='Site B', slug='site-b')
+
+        manufacturer = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        device_type = DeviceType.objects.create(
+            manufacturer=manufacturer, model='Device Type 1', slug='device-type-1'
+        )
+        device_role = DeviceRole.objects.create(
+            name='Device Role 1', slug='device-role-1', color='ff0000'
+        )
+
+        # Create Rack1 in Site A
+        rack1 = Rack.objects.create(site=site_a, name='Rack 1')
+
+        # Create Device1 in Rack1
+        device1 = Device.objects.create(site=site_a, rack=rack1, device_type=device_type, device_role=device_role)
+
+        # Move Rack1 to Site B
+        rack1.site = site_b
+        rack1.save()
+
+        # Check that Device1 is now assigned to Site B
+        self.assertEqual(Device.objects.get(pk=device1.pk).site, site_b)
+
 
 class DeviceTestCase(TestCase):