From 36cf40f25cf8ad6b231cba3a2b8d387cced0abd6 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Wed, 24 Jun 2020 15:29:25 -0400
Subject: [PATCH] Enable CSV import for virtual chassis

---
 netbox/dcim/forms.py               | 13 +++++++++++++
 netbox/dcim/tests/test_views.py    | 13 +++++++++++--
 netbox/dcim/urls.py                |  1 +
 netbox/dcim/views.py               |  7 +++++++
 netbox/templates/inc/nav_menu.html |  1 +
 5 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/netbox/dcim/forms.py b/netbox/dcim/forms.py
index 6f51160fa..02b6eaa6e 100644
--- a/netbox/dcim/forms.py
+++ b/netbox/dcim/forms.py
@@ -4306,6 +4306,19 @@ class VirtualChassisBulkEditForm(BootstrapMixin, AddRemoveTagsForm, BulkEditForm
         nullable_fields = ['domain']
 
 
+class VirtualChassisCSVForm(CSVModelForm):
+    master = CSVModelChoiceField(
+        queryset=Device.objects.all(),
+        to_field_name='name',
+        required=False,
+        help_text='Master device'
+    )
+
+    class Meta:
+        model = VirtualChassis
+        fields = VirtualChassis.csv_headers
+
+
 class VirtualChassisFilterForm(BootstrapMixin, CustomFieldFilterForm):
     model = VirtualChassis
     q = forms.CharField(
diff --git a/netbox/dcim/tests/test_views.py b/netbox/dcim/tests/test_views.py
index 23e65eb05..5c5e46853 100644
--- a/netbox/dcim/tests/test_views.py
+++ b/netbox/dcim/tests/test_views.py
@@ -1578,7 +1578,6 @@ class VirtualChassisTestCase(ViewTestCases.PrimaryObjectViewTestCase):
             name='Device Role', slug='device-role-1'
         )
 
-        # Create 9 member Devices
         devices = (
             Device(device_type=device_type, device_role=device_role, name='Device 1', site=site),
             Device(device_type=device_type, device_role=device_role, name='Device 2', site=site),
@@ -1589,10 +1588,13 @@ class VirtualChassisTestCase(ViewTestCases.PrimaryObjectViewTestCase):
             Device(device_type=device_type, device_role=device_role, name='Device 7', site=site),
             Device(device_type=device_type, device_role=device_role, name='Device 8', site=site),
             Device(device_type=device_type, device_role=device_role, name='Device 9', site=site),
+            Device(device_type=device_type, device_role=device_role, name='Device 10', site=site),
+            Device(device_type=device_type, device_role=device_role, name='Device 11', site=site),
+            Device(device_type=device_type, device_role=device_role, name='Device 12', site=site),
         )
         Device.objects.bulk_create(devices)
 
-        # Create three VirtualChassis with two members each
+        # Create three VirtualChassis with three members each
         vc1 = VirtualChassis.objects.create(name='VC1', master=devices[0], domain='domain-1')
         Device.objects.filter(pk=devices[0].pk).update(virtual_chassis=vc1, vc_position=1)
         Device.objects.filter(pk=devices[1].pk).update(virtual_chassis=vc1, vc_position=2)
@@ -1616,6 +1618,13 @@ class VirtualChassisTestCase(ViewTestCases.PrimaryObjectViewTestCase):
             'form-MAX_NUM_FORMS': 1000,
         }
 
+        cls.csv_data = (
+            "name,domain,master",
+            "VC4,Domain 4,Device 10",
+            "VC5,Domain 5,Device 11",
+            "VC6,Domain 6,Device 12",
+        )
+
         cls.bulk_edit_data = {
             'domain': 'domain-x',
         }
diff --git a/netbox/dcim/urls.py b/netbox/dcim/urls.py
index a0d6bdc92..347ac7064 100644
--- a/netbox/dcim/urls.py
+++ b/netbox/dcim/urls.py
@@ -321,6 +321,7 @@ urlpatterns = [
     # Virtual chassis
     path('virtual-chassis/', views.VirtualChassisListView.as_view(), name='virtualchassis_list'),
     path('virtual-chassis/add/', views.VirtualChassisCreateView.as_view(), name='virtualchassis_add'),
+    path('virtual-chassis/import/', views.VirtualChassisBulkImportView.as_view(), name='virtualchassis_import'),
     path('virtual-chassis/edit/', views.VirtualChassisBulkEditView.as_view(), name='virtualchassis_bulk_edit'),
     path('virtual-chassis/delete/', views.VirtualChassisBulkDeleteView.as_view(), name='virtualchassis_bulk_delete'),
     path('virtual-chassis/<int:pk>/', views.VirtualChassisView.as_view(), name='virtualchassis'),
diff --git a/netbox/dcim/views.py b/netbox/dcim/views.py
index f6fe7cf74..840b9890f 100644
--- a/netbox/dcim/views.py
+++ b/netbox/dcim/views.py
@@ -2304,6 +2304,13 @@ class VirtualChassisRemoveMemberView(ObjectPermissionRequiredMixin, GetReturnURL
         })
 
 
+class VirtualChassisBulkImportView(BulkImportView):
+    queryset = VirtualChassis.objects.all()
+    model_form = forms.VirtualChassisCSVForm
+    table = tables.VirtualChassisTable
+    default_return_url = 'dcim:virtualchassis_list'
+
+
 class VirtualChassisBulkEditView(BulkEditView):
     queryset = VirtualChassis.objects.all()
     filterset = filters.VirtualChassisFilterSet
diff --git a/netbox/templates/inc/nav_menu.html b/netbox/templates/inc/nav_menu.html
index 6aa80f910..f22baf7cc 100644
--- a/netbox/templates/inc/nav_menu.html
+++ b/netbox/templates/inc/nav_menu.html
@@ -147,6 +147,7 @@
                             {% if perms.dcim.add_virtualchassis %}
                                 <div class="buttons pull-right">
                                     <a href="{% url 'dcim:virtualchassis_add' %}" class="btn btn-xs btn-success" title="Add"><i class="fa fa-plus"></i></a>
+                                    <a href="{% url 'dcim:virtualchassis_import' %}" class="btn btn-xs btn-info" title="Import"><i class="fa fa-download"></i></a>
                                 </div>
                             {% endif %}
                             <a href="{% url 'dcim:virtualchassis_list' %}">Virtual Chassis</a>