Просмотр исходного кода

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 дней назад
Родитель
Сommit
80cc7e0d91

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

@@ -1,5 +1,3 @@
-from collections import defaultdict
-
 import jsonschema
 from django.conf import settings
 from django.core.validators import ValidationError
@@ -8,7 +6,6 @@ from django.urls import reverse
 from django.utils.translation import gettext_lazy as _
 from jsonschema.exceptions import ValidationError as JSONValidationError
 
-from core.models import ObjectType
 from extras.models.mixins import RenderTemplateMixin
 from extras.querysets import ConfigContextQuerySet
 from netbox.models import ChangeLoggedModel, PrimaryModel
@@ -302,17 +299,3 @@ class ConfigTemplate(
         """
         self.template_code = self.data_file.data_as_string
     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 os
 import sys
+from collections import defaultdict
 
 from django.core.files.storage import storages
 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.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.utils import filename_from_model, filename_from_object
 from utilities.jinja2 import render_jinja2
@@ -120,9 +122,17 @@ class RenderTemplateMixin(models.Model):
         abstract = True
 
     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):
         """

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

@@ -458,14 +458,8 @@ class ExportTemplate(
     sync_data.alters_data = True
 
     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
 
 

+ 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 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 utilities.exceptions import AbortRequest
 from virtualization.models import Cluster, ClusterGroup, ClusterType, VirtualMachine
@@ -808,3 +816,36 @@ class ConfigTemplateTest(TestCase):
                 object_id=config_template.pk
             )
             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)