From 8d8c58df549e891d2923c0bab65ad636636942e8 Mon Sep 17 00:00:00 2001
From: ryanmerolle <ryan.merolle@gmail.com>
Date: Wed, 30 Dec 2020 19:09:08 -0500
Subject: [PATCH] optional assoc to cluster & circuit startup_script

---
 initializers/circuits.yml       |  1 +
 initializers/clusters.yml       |  1 +
 startup_scripts/170_clusters.py |  4 +++-
 startup_scripts/300_circuits.py | 12 ++++++++++++
 4 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/initializers/circuits.yml b/initializers/circuits.yml
index 93b4659..d47b06f 100644
--- a/initializers/circuits.yml
+++ b/initializers/circuits.yml
@@ -1,6 +1,7 @@
 # - cid: Circuit_ID-1
 #   provider: Provider1
 #   type: Internet
+#   tenant: tenant1
 # - cid: Circuit_ID-2
 #   provider: Provider2
 #   type: MPLS
diff --git a/initializers/clusters.yml b/initializers/clusters.yml
index 692be6c..4e15477 100644
--- a/initializers/clusters.yml
+++ b/initializers/clusters.yml
@@ -1,6 +1,7 @@
 # - name: cluster1
 #   type: Hyper-V
 #   group: Group 1
+#   tenant: tenant1
 # - name: cluster2
 #   type: Hyper-V
 #   site: SING 1
diff --git a/startup_scripts/170_clusters.py b/startup_scripts/170_clusters.py
index ffd965e..85e462f 100644
--- a/startup_scripts/170_clusters.py
+++ b/startup_scripts/170_clusters.py
@@ -3,6 +3,7 @@ import sys
 from dcim.models import Site
 from startup_script_utils import *
 from virtualization.models import Cluster, ClusterType, ClusterGroup
+from tenancy.models import Tenant
 
 clusters = load_yaml('/opt/netbox/initializers/clusters.yml')
 
@@ -15,7 +16,8 @@ required_assocs = {
 
 optional_assocs = {
   'site': (Site, 'name'),
-  'group': (ClusterGroup, 'name')
+  'group': (ClusterGroup, 'name'),
+  'tenant': (Tenant, 'name')
 }
 
 for params in clusters:
diff --git a/startup_scripts/300_circuits.py b/startup_scripts/300_circuits.py
index 456c792..d53272c 100644
--- a/startup_scripts/300_circuits.py
+++ b/startup_scripts/300_circuits.py
@@ -1,4 +1,5 @@
 from circuits.models import Circuit, Provider, CircuitType
+from tenancy.models import Tenant
 from startup_script_utils import *
 import sys
 
@@ -12,6 +13,10 @@ required_assocs = {
   'type': (CircuitType, 'name')
 }
 
+optional_assocs = {
+  'tenant': (Tenant, 'name')
+}
+
 for params in circuits:
   custom_field_data = pop_custom_fields(params)
 
@@ -21,6 +26,13 @@ for params in circuits:
 
     params[assoc] = model.objects.get(**query)
 
+  for assoc, details in optional_assocs.items():
+    if assoc in params:
+      model, field = details
+      query = { field: params.pop(assoc) }
+
+      params[assoc] = model.objects.get(**query)
+
   circuit, created = Circuit.objects.get_or_create(**params)
 
   if created: