From df95115e2e559afc5866390d6d06f66ae7db8ae0 Mon Sep 17 00:00:00 2001
From: jeremystretch <jstretch@ns1.com>
Date: Thu, 27 Jan 2022 13:37:24 -0500
Subject: [PATCH] Refactor plugins registry

---
 netbox/extras/plugins/__init__.py     | 16 +++++++++-------
 netbox/extras/templatetags/plugins.py |  2 +-
 netbox/extras/tests/test_plugins.py   | 10 +++++-----
 netbox/netbox/navigation_menu.py      |  4 ++--
 netbox/netbox/preferences.py          |  4 ++--
 5 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/netbox/extras/plugins/__init__.py b/netbox/extras/plugins/__init__.py
index 5b02b5ab7..f297d43e1 100644
--- a/netbox/extras/plugins/__init__.py
+++ b/netbox/extras/plugins/__init__.py
@@ -12,10 +12,12 @@ from utilities.choices import ButtonColorChoices
 from extras.plugins.utils import import_object
 
 
-# Initialize plugin registry stores
-registry['plugin_template_extensions'] = collections.defaultdict(list)
-registry['plugin_menu_items'] = {}
-registry['plugin_preferences'] = {}
+# Initialize plugin registry
+registry['plugins'] = {
+    'template_extensions': collections.defaultdict(list),
+    'menu_items': {},
+    'preferences': {},
+}
 
 
 #
@@ -184,7 +186,7 @@ def register_template_extensions(class_list):
         if template_extension.model is None:
             raise TypeError(f"PluginTemplateExtension class {template_extension} does not define a valid model!")
 
-        registry['plugin_template_extensions'][template_extension.model].append(template_extension)
+        registry['plugins']['template_extensions'][template_extension.model].append(template_extension)
 
 
 #
@@ -249,7 +251,7 @@ def register_menu_items(section_name, class_list):
             if not isinstance(button, PluginMenuButton):
                 raise TypeError(f"{button} must be an instance of extras.plugins.PluginMenuButton")
 
-    registry['plugin_menu_items'][section_name] = class_list
+    registry['plugins']['menu_items'][section_name] = class_list
 
 
 #
@@ -260,4 +262,4 @@ def register_user_preferences(plugin_name, preferences):
     """
     Register a list of user preferences defined by a plugin.
     """
-    registry['plugin_preferences'][plugin_name] = preferences
+    registry['plugins']['preferences'][plugin_name] = preferences
diff --git a/netbox/extras/templatetags/plugins.py b/netbox/extras/templatetags/plugins.py
index 200b78e34..df3024a16 100644
--- a/netbox/extras/templatetags/plugins.py
+++ b/netbox/extras/templatetags/plugins.py
@@ -23,7 +23,7 @@ def _get_registered_content(obj, method, template_context):
     }
 
     model_name = obj._meta.label_lower
-    template_extensions = registry['plugin_template_extensions'].get(model_name, [])
+    template_extensions = registry['plugins']['template_extensions'].get(model_name, [])
     for template_extension in template_extensions:
 
         # If the class has not overridden the specified method, we can skip it (because we know it
diff --git a/netbox/extras/tests/test_plugins.py b/netbox/extras/tests/test_plugins.py
index 4bea9933e..608fc58dc 100644
--- a/netbox/extras/tests/test_plugins.py
+++ b/netbox/extras/tests/test_plugins.py
@@ -61,8 +61,8 @@ class PluginTest(TestCase):
         """
         Check that plugin MenuItems and MenuButtons are registered.
         """
-        self.assertIn('Dummy plugin', registry['plugin_menu_items'])
-        menu_items = registry['plugin_menu_items']['Dummy plugin']
+        self.assertIn('Dummy plugin', registry['plugins']['menu_items'])
+        menu_items = registry['plugins']['menu_items']['Dummy plugin']
         self.assertEqual(len(menu_items), 2)
         self.assertEqual(len(menu_items[0].buttons), 2)
 
@@ -72,14 +72,14 @@ class PluginTest(TestCase):
         """
         from extras.tests.dummy_plugin.template_content import SiteContent
 
-        self.assertIn(SiteContent, registry['plugin_template_extensions']['dcim.site'])
+        self.assertIn(SiteContent, registry['plugins']['template_extensions']['dcim.site'])
 
     def test_user_preferences(self):
         """
         Check that plugin UserPreferences are registered.
         """
-        self.assertIn('dummy_plugin', registry['plugin_preferences'])
-        user_preferences = registry['plugin_preferences']['dummy_plugin']
+        self.assertIn('dummy_plugin', registry['plugins']['preferences'])
+        user_preferences = registry['plugins']['preferences']['dummy_plugin']
         self.assertEqual(type(user_preferences), dict)
         self.assertEqual(list(user_preferences.keys()), ['pref1', 'pref2'])
 
diff --git a/netbox/netbox/navigation_menu.py b/netbox/netbox/navigation_menu.py
index 85d86a47a..9a55c263e 100644
--- a/netbox/netbox/navigation_menu.py
+++ b/netbox/netbox/navigation_menu.py
@@ -390,10 +390,10 @@ MENUS = [
 # Add plugin menus
 #
 
-if registry['plugin_menu_items']:
+if registry['plugins']['menu_items']:
     plugin_menu_groups = []
 
-    for plugin_name, items in registry['plugin_menu_items'].items():
+    for plugin_name, items in registry['plugins']['menu_items'].items():
         plugin_menu_groups.append(
             MenuGroup(
                 label=plugin_name,
diff --git a/netbox/netbox/preferences.py b/netbox/netbox/preferences.py
index aec8bc752..6bf56b562 100644
--- a/netbox/netbox/preferences.py
+++ b/netbox/netbox/preferences.py
@@ -49,10 +49,10 @@ PREFERENCES = {
 }
 
 # Register plugin preferences
-if registry['plugin_preferences']:
+if registry['plugins']['preferences']:
     plugin_preferences = {}
 
-    for plugin_name, preferences in registry['plugin_preferences'].items():
+    for plugin_name, preferences in registry['plugins']['preferences'].items():
         for name, userpreference in preferences.items():
             PREFERENCES[f'plugins.{plugin_name}.{name}'] = userpreference