Repository: nautobot/nautobot

Issue:
2.0: Delaring Any Model with `@extras_features("relationships")` Empties Those from #826
<!--
    NOTE: IF YOUR ISSUE DOES NOT FOLLOW THIS TEMPLATE, IT WILL BE CLOSED.

    This form is only for reporting reproducible bugs. If you need assistance
    with Nautobot installation, or if you have a general question, please start a
    discussion instead: https://github.com/nautobot/nautobot/discussions

    Please describe the environment in which you are running Nautobot. Be sure
    that you are running an unmodified instance of the latest stable release
    before submitting a bug report, and that any plugins have been disabled.
-->
### Environment
* Nautobot version (Docker tag too if applicable): `next` or `v2.0.0-alpha.3`
* Python version: 3.X
* Database platform, version: Any
* Middleware(s): N/A

<!--
    Describe in detail the exact steps that someone else can take to reproduce
    this bug using the current stable release of Nautobot. Begin with the
    creation of any necessary database objects and call out every operation
    being performed explicitly. If reporting a bug in the REST API, be sure to
    reconstruct the raw HTTP request(s) being made: Don't rely on a client
    library such as pynautobot.
-->

Same applies for Custom Fields.


### Steps to Reproduce
1. Clone the repo
2. Add `@extras_features("relationships")` on any model inheriting from `PrimaryModel`
3. Run Tests/Edit or Create a Relationship

<!-- What did you expect to happen? -->
### Expected Behavior
Have access to all models that inherit from `PrimaryModel` to associate a Relationship to.


<!-- What happened instead? -->
### Observed Behavior
Only have the decorated model(s).



Discussion:

Resolving patch:
diff --git a/examples/example_plugin/example_plugin/models.py b/examples/example_plugin/example_plugin/models.py
--- a/examples/example_plugin/example_plugin/models.py
+++ b/examples/example_plugin/example_plugin/models.py
@@ -40,6 +40,7 @@ def to_csv(self):
     "export_templates",
     # "graphql", Not specified here as we have a custom type for this model, see example_plugin.graphql.types
     "webhooks",
+    "relationships",  # Defined here to ensure no clobbering: https://github.com/nautobot/nautobot/issues/3592
 )
 class AnotherExampleModel(OrganizationalModel):
     name = models.CharField(max_length=20)
diff --git a/nautobot/extras/utils.py b/nautobot/extras/utils.py
--- a/nautobot/extras/utils.py
+++ b/nautobot/extras/utils.py
@@ -136,11 +136,8 @@ def get_query(self):
         # initialization of the application, before `ExtrasConfig.ready` is called.
         # Calling `populate_model_features_registry` in `ExtrasConfig.ready` would lead to an outdated `model_features`
         # `registry` record being used by `FeatureQuery`.
-        # As the `populate_model_features_registry` can be resource-intensive. The check the conditional check is used to
-        # avoid calling the function multiple times and optimize the performance of the application.
-        # TODO(timizuo): Provide a better solution for this; check https://github.com/nautobot/nautobot/pull/3360/files#r1131373797 comment.
-        if not registry["model_features"].get("relationships"):
-            populate_model_features_registry()
+
+        populate_model_features_registry()
         query = Q()
         for app_label, models in registry["model_features"][self.feature].items():
             query |= Q(app_label=app_label, model__in=models)
@@ -223,7 +220,7 @@ def wrapper(model_class):
     return wrapper
 
 
-def populate_model_features_registry():
+def populate_model_features_registry(refresh=False):
     """
     Populate the registry model features with new apps.
 
@@ -243,6 +240,9 @@ def populate_model_features_registry():
     - For each dictionary in lookup_confs, calls lookup_by_field() function to look for all models that have fields with the names given in the dictionary.
     - Groups the results by app and updates the registry model features for each app.
     """
+    if registry.get("populate_model_features_registry_called", False) and not refresh:
+        return
+
     RelationshipAssociation = apps.get_model(app_label="extras", model_name="relationshipassociation")
 
     lookup_confs = [
@@ -267,6 +267,9 @@ def populate_model_features_registry():
         feature_name = lookup_conf["feature_name"]
         registry["model_features"][feature_name] = registry_items
 
+    if not registry.get("populate_model_features_registry_called", False):
+        registry["populate_model_features_registry_called"] = True
+
 
 def generate_signature(request_body, secret):
     """


Test patch:
diff --git a/nautobot/extras/tests/test_utils.py b/nautobot/extras/tests/test_utils.py
--- a/nautobot/extras/tests/test_utils.py
+++ b/nautobot/extras/tests/test_utils.py
@@ -1,7 +1,8 @@
 from unittest import mock
 
 from nautobot.core.testing import TestCase
-from nautobot.extras.utils import get_celery_queues, get_worker_count
+from nautobot.extras.registry import registry
+from nautobot.extras.utils import get_celery_queues, get_worker_count, populate_model_features_registry
 
 
 class UtilsTestCase(TestCase):
@@ -47,3 +48,32 @@ def test_get_worker_count(self, mock_get_celery_queues):
             self.assertEqual(get_worker_count(queue="bulk"), 4)
         with self.subTest("Empty queue"):
             self.assertEqual(get_worker_count(queue="nobody"), 0)
+
+    def test_populate_model_features_registry(self):
+        original_custom_fields_registry = registry["model_features"]["custom_fields"].copy()
+
+        self.assertIn(
+            "circuit", registry["model_features"]["custom_fields"]["circuits"], "Registry should already be populated"
+        )
+
+        populate_model_features_registry()
+        self.assertDictEqual(
+            registry["model_features"]["custom_fields"],
+            original_custom_fields_registry,
+            "Registry should not be modified if refresh flag not set",
+        )
+
+        registry["model_features"]["custom_fields"].pop("circuits")
+        self.assertNotEqual(
+            registry["model_features"]["custom_fields"],
+            original_custom_fields_registry,
+            "Registry should be successfully modified",
+        )
+
+        # Make sure we return to the original state
+        populate_model_features_registry(refresh=True)
+        self.assertDictEqual(
+            registry["model_features"]["custom_fields"],
+            original_custom_fields_registry,
+            "Registry should be restored to original state",
+        )