netbox/netbox/extras/tests/test_utils.py

from types import SimpleNamespace

from django.contrib.contenttypes.models import ContentType
from django.test import TestCase

from extras.models import ExportTemplate
from extras.utils import filename_from_model, image_upload
from tenancy.models import ContactGroup, TenantGroup
from wireless.models import WirelessLANGroup


class FilenameFromModelTests(TestCase):
    def test_expected_output(self):
        cases = (
            (ExportTemplate, 'netbox_export_templates'),
            (ContactGroup, 'netbox_contact_groups'),
            (TenantGroup, 'netbox_tenant_groups'),
            (WirelessLANGroup, 'netbox_wireless_lan_groups'),
        )

        for model, expected in cases:
            self.assertEqual(filename_from_model(model), expected)


class ImageUploadTests(TestCase):
    @classmethod
    def setUpTestData(cls):
        # We only need a ContentType with model="rack" for the prefix;
        # this doesn't require creating a Rack object.
        cls.ct_rack = ContentType.objects.get_by_natural_key('dcim', 'rack')

    def _stub_instance(self, object_id=12, name=None):
        """
        Creates a minimal stub for use with the `image_upload()` function.

        This method generates an instance of `SimpleNamespace` containing a set
        of attributes required to simulate the expected input for the
        `image_upload()` method.
        It is designed to simplify testing or processing by providing a
        lightweight representation of an object.
        """
        return SimpleNamespace(object_type=self.ct_rack, object_id=object_id, name=name)

    def _second_segment(self, path: str):
        """
        Extracts and returns the portion of the input string after the
        first '/' character.
        """
        return path.split('/', 1)[1]

    def test_windows_fake_path_and_extension_lowercased(self):
        """
        Tests handling of a Windows file path with a fake directory and extension.
        """
        inst = self._stub_instance(name=None)
        path = image_upload(inst, r'C:\fake_path\MyPhoto.JPG')
        # Base directory and single-level path
        seg2 = self._second_segment(path)
        self.assertTrue(path.startswith('image-attachments/rack_12_'))
        self.assertNotIn('/', seg2, 'should not create nested directories')
        # Extension from the uploaded file, lowercased
        self.assertTrue(seg2.endswith('.jpg'))

    def test_name_with_slashes_is_flattened_no_subdirectories(self):
        """
        Tests that a name with slashes is flattened and does not
        create subdirectories.
        """
        inst = self._stub_instance(name='5/31/23')
        path = image_upload(inst, 'image.png')
        seg2 = self._second_segment(path)
        self.assertTrue(seg2.startswith('rack_12_'))
        self.assertNotIn('/', seg2)
        self.assertNotIn('\\', seg2)
        self.assertTrue(seg2.endswith('.png'))

    def test_name_with_backslashes_is_flattened_no_subdirectories(self):
        """
        Tests that a name including backslashes is correctly flattened
        into a single directory name without creating subdirectories.
        """
        inst = self._stub_instance(name=r'5\31\23')
        path = image_upload(inst, 'image_name.png')

        seg2 = self._second_segment(path)
        self.assertTrue(seg2.startswith('rack_12_'))
        self.assertNotIn('/', seg2)
        self.assertNotIn('\\', seg2)
        self.assertTrue(seg2.endswith('.png'))

    def test_prefix_format_is_as_expected(self):
        """
        Tests the output path format generated by the `image_upload` function.
        """
        inst = self._stub_instance(object_id=99, name='label')
        path = image_upload(inst, 'a.webp')
        # The second segment must begin with "rack_99_"
        seg2 = self._second_segment(path)
        self.assertTrue(seg2.startswith('rack_99_'))
        self.assertTrue(seg2.endswith('.webp'))

    def test_unsupported_file_extension(self):
        """
        Test that when the file extension is not allowed, the extension
        is omitted.
        """
        inst = self._stub_instance(name='test')
        path = image_upload(inst, 'document.txt')

        seg2 = self._second_segment(path)
        self.assertTrue(seg2.startswith('rack_12_test'))
        self.assertFalse(seg2.endswith('.txt'))
        # When not allowed, no extension should be appended
        self.assertNotRegex(seg2, r'\.txt$')

    def test_instance_name_with_whitespace_and_special_chars(self):
        """
        Test that an instance name with leading/trailing whitespace and
        special characters is sanitized properly.
        """
        # Suppose the instance name has surrounding whitespace and
        # extra slashes.
        inst = self._stub_instance(name='  my/complex\\name  ')
        path = image_upload(inst, 'irrelevant.png')

        # The output should be flattened and sanitized.
        # We expect the name to be transformed into a valid filename without
        # path separators.
        seg2 = self._second_segment(path)
        self.assertNotIn(' ', seg2)
        self.assertNotIn('/', seg2)
        self.assertNotIn('\\', seg2)
        self.assertTrue(seg2.endswith('.png'))

    def test_separator_variants_with_subTest(self):
        """
        Tests that both forward slash and backslash in file paths are
        handled consistently by the `image_upload` function and
        processed into a sanitized uniform format.
        """
        for name in ['2025/09/12', r'2025\09\12']:
            with self.subTest(name=name):
                inst = self._stub_instance(name=name)
                path = image_upload(inst, 'x.jpeg')
                seg2 = self._second_segment(path)
                self.assertTrue(seg2.startswith('rack_12_'))
                self.assertNotIn('/', seg2)
                self.assertNotIn('\\', seg2)
                self.assertTrue(seg2.endswith('.jpeg') or seg2.endswith('.jpg'))

    def test_fallback_on_suspicious_file_operation(self):
        """
        Test that when default_storage.get_valid_name() raises a
        SuspiciousFileOperation, the fallback default is used.
        """
        inst = self._stub_instance(name=' ')
        path = image_upload(inst, 'sample.png')
        # Expect the fallback name 'unnamed' to be used.
        self.assertIn('unnamed', path)
        self.assertTrue(path.startswith('image-attachments/rack_12_'))