Răsfoiți Sursa

Closes #21157: Add public models to export template context

Move shared get_context() logic from ConfigTemplate into
RenderTemplateMixin so ExportTemplate also gets access to all
public model classes. This enables export templates to perform
cross-model lookups (e.g. resolving parent Prefix from IPAddress).
Jason Novinger 2 zile în urmă
părinte
comite
80cc7e0d91

+ 0 - 17
netbox/extras/models/configs.py

@@ -1,5 +1,3 @@
-from collections import defaultdict
-
 import jsonschema
 import jsonschema
 from django.conf import settings
 from django.conf import settings
 from django.core.validators import ValidationError
 from django.core.validators import ValidationError
@@ -8,7 +6,6 @@ from django.urls import reverse
 from django.utils.translation import gettext_lazy as _
 from django.utils.translation import gettext_lazy as _
 from jsonschema.exceptions import ValidationError as JSONValidationError
 from jsonschema.exceptions import ValidationError as JSONValidationError
 
 
-from core.models import ObjectType
 from extras.models.mixins import RenderTemplateMixin
 from extras.models.mixins import RenderTemplateMixin
 from extras.querysets import ConfigContextQuerySet
 from extras.querysets import ConfigContextQuerySet
 from netbox.models import ChangeLoggedModel, PrimaryModel
 from netbox.models import ChangeLoggedModel, PrimaryModel
@@ -302,17 +299,3 @@ class ConfigTemplate(
         """
         """
         self.template_code = self.data_file.data_as_string
         self.template_code = self.data_file.data_as_string
     sync_data.alters_data = True
     sync_data.alters_data = True
-
-    def get_context(self, context=None, queryset=None):
-        _context = defaultdict(dict)
-
-        # Populate all public models for reference within the template
-        for object_type in ObjectType.objects.public():
-            if model := object_type.model_class():
-                _context[object_type.app_label][model.__name__] = model
-
-        # Apply the provided context data, if any
-        if context is not None:
-            _context.update(context)
-
-        return _context

+ 13 - 3
netbox/extras/models/mixins.py

@@ -2,6 +2,7 @@ import importlib.abc
 import importlib.util
 import importlib.util
 import os
 import os
 import sys
 import sys
+from collections import defaultdict
 
 
 from django.core.files.storage import storages
 from django.core.files.storage import storages
 from django.db import models
 from django.db import models
@@ -9,6 +10,7 @@ from django.http import HttpResponse
 from django.utils.module_loading import import_string
 from django.utils.module_loading import import_string
 from django.utils.translation import gettext_lazy as _
 from django.utils.translation import gettext_lazy as _
 
 
+from core.models import ObjectType
 from extras.constants import DEFAULT_MIME_TYPE, JINJA_ENV_PARAMS_WITH_PATH_IMPORT
 from extras.constants import DEFAULT_MIME_TYPE, JINJA_ENV_PARAMS_WITH_PATH_IMPORT
 from extras.utils import filename_from_model, filename_from_object
 from extras.utils import filename_from_model, filename_from_object
 from utilities.jinja2 import render_jinja2
 from utilities.jinja2 import render_jinja2
@@ -120,9 +122,17 @@ class RenderTemplateMixin(models.Model):
         abstract = True
         abstract = True
 
 
     def get_context(self, context=None, queryset=None):
     def get_context(self, context=None, queryset=None):
-        raise NotImplementedError(_("{class_name} must implement a get_context() method.").format(
-            class_name=self.__class__
-        ))
+        _context = defaultdict(dict)
+
+        # Populate all public models for reference within the template
+        for object_type in ObjectType.objects.public():
+            if model := object_type.model_class():
+                _context[object_type.app_label][model.__name__] = model
+
+        if context is not None:
+            _context.update(context)
+
+        return _context
 
 
     def get_environment_params(self):
     def get_environment_params(self):
         """
         """

+ 2 - 8
netbox/extras/models/models.py

@@ -458,14 +458,8 @@ class ExportTemplate(
     sync_data.alters_data = True
     sync_data.alters_data = True
 
 
     def get_context(self, context=None, queryset=None):
     def get_context(self, context=None, queryset=None):
-        _context = {
-            'queryset': queryset,
-        }
-
-        # Apply the provided context data, if any
-        if context is not None:
-            _context.update(context)
-
+        _context = super().get_context(context=context, queryset=queryset)
+        _context['queryset'] = queryset
         return _context
         return _context
 
 
 
 

+ 42 - 1
netbox/extras/tests/test_models.py

@@ -8,7 +8,15 @@ from django.test import TestCase, tag
 
 
 from core.models import AutoSyncRecord, DataSource, ObjectType
 from core.models import AutoSyncRecord, DataSource, ObjectType
 from dcim.models import Device, DeviceRole, DeviceType, Location, Manufacturer, Platform, Region, Site, SiteGroup
 from dcim.models import Device, DeviceRole, DeviceType, Location, Manufacturer, Platform, Region, Site, SiteGroup
-from extras.models import ConfigContext, ConfigContextProfile, ConfigTemplate, ImageAttachment, Tag, TaggedItem
+from extras.models import (
+    ConfigContext,
+    ConfigContextProfile,
+    ConfigTemplate,
+    ExportTemplate,
+    ImageAttachment,
+    Tag,
+    TaggedItem,
+)
 from tenancy.models import Tenant, TenantGroup
 from tenancy.models import Tenant, TenantGroup
 from utilities.exceptions import AbortRequest
 from utilities.exceptions import AbortRequest
 from virtualization.models import Cluster, ClusterGroup, ClusterType, VirtualMachine
 from virtualization.models import Cluster, ClusterGroup, ClusterType, VirtualMachine
@@ -808,3 +816,36 @@ class ConfigTemplateTest(TestCase):
                 object_id=config_template.pk
                 object_id=config_template.pk
             )
             )
             self.assertEqual(autosync_records.count(), 0, "AutoSyncRecord should be deleted after detaching")
             self.assertEqual(autosync_records.count(), 0, "AutoSyncRecord should be deleted after detaching")
+
+
+class ExportTemplateContextTest(TestCase):
+    """
+    Tests for ExportTemplate.get_context() including public model population.
+    """
+
+    def test_get_context_includes_public_models(self):
+        et = ExportTemplate(name='test', template_code='test')
+        ctx = et.get_context()
+
+        self.assertIs(ctx['dcim']['Site'], Site)
+        self.assertIs(ctx['dcim']['Device'], Device)
+
+    def test_get_context_includes_queryset(self):
+        et = ExportTemplate(name='test', template_code='test')
+        qs = Site.objects.all()
+        ctx = et.get_context(queryset=qs)
+
+        self.assertIs(ctx['queryset'], qs)
+
+    def test_get_context_applies_extra_context(self):
+        et = ExportTemplate(name='test', template_code='test')
+        ctx = et.get_context(context={'custom_key': 'custom_value'})
+
+        self.assertEqual(ctx['custom_key'], 'custom_value')
+        self.assertIs(ctx['dcim']['Site'], Site)
+
+    def test_config_template_get_context_includes_public_models(self):
+        ct = ConfigTemplate(name='test', template_code='test')
+        ctx = ct.get_context()
+
+        self.assertIs(ctx['dcim']['Site'], Site)