netbox/netbox/utilities/tests/test_permissions.py

from django.test import TestCase

from core.models import ObjectType
from dcim.models import Device, Site
from netbox.registry import registry
from users.forms.model_forms import ObjectPermissionForm
from users.models import ObjectPermission
from utilities.permissions import ModelAction, register_model_actions
from virtualization.models import VirtualMachine


class ModelActionTest(TestCase):

    def test_hash(self):
        action1 = ModelAction(name='sync')
        action2 = ModelAction(name='sync', help_text='Different help')
        self.assertEqual(hash(action1), hash(action2))

    def test_equality_with_model_action(self):
        action1 = ModelAction(name='sync')
        action2 = ModelAction(name='sync', help_text='Different help')
        action3 = ModelAction(name='merge')
        self.assertEqual(action1, action2)
        self.assertNotEqual(action1, action3)

    def test_equality_with_string(self):
        action = ModelAction(name='sync')
        self.assertEqual(action, 'sync')
        self.assertNotEqual(action, 'merge')

    def test_usable_in_set(self):
        action1 = ModelAction(name='sync')
        action2 = ModelAction(name='sync', help_text='Different')
        action3 = ModelAction(name='merge')
        actions = {action1, action2, action3}
        self.assertEqual(len(actions), 2)


class RegisterModelActionsTest(TestCase):

    def setUp(self):
        self._original_actions = dict(registry['model_actions'])
        registry['model_actions'].clear()

    def tearDown(self):
        registry['model_actions'].clear()
        registry['model_actions'].update(self._original_actions)

    def test_register_model_action_objects(self):
        register_model_actions(Site, [
            ModelAction('test_action', help_text='Test help'),
        ])
        actions = registry['model_actions']['dcim.site']
        self.assertEqual(len(actions), 1)
        self.assertEqual(actions[0].name, 'test_action')
        self.assertEqual(actions[0].help_text, 'Test help')

    def test_register_string_actions(self):
        register_model_actions(Site, ['action1', 'action2'])
        actions = registry['model_actions']['dcim.site']
        self.assertEqual(len(actions), 2)
        self.assertIsInstance(actions[0], ModelAction)
        self.assertEqual(actions[0].name, 'action1')
        self.assertEqual(actions[1].name, 'action2')

    def test_register_mixed_actions(self):
        register_model_actions(Site, [
            ModelAction('with_help', help_text='Has help'),
            'without_help',
        ])
        actions = registry['model_actions']['dcim.site']
        self.assertEqual(len(actions), 2)
        self.assertEqual(actions[0].help_text, 'Has help')
        self.assertEqual(actions[1].help_text, '')

    def test_multiple_registrations_append(self):
        register_model_actions(Site, [ModelAction('first')])
        register_model_actions(Site, [ModelAction('second')])
        actions = registry['model_actions']['dcim.site']
        self.assertEqual(len(actions), 2)
        self.assertEqual(actions[0].name, 'first')
        self.assertEqual(actions[1].name, 'second')

    def test_duplicate_registration_ignored(self):
        register_model_actions(Site, [ModelAction('sync')])
        register_model_actions(Site, [ModelAction('sync', help_text='Different help')])
        actions = registry['model_actions']['dcim.site']
        self.assertEqual(len(actions), 1)

    def test_reserved_action_rejected(self):
        for action_name in ('view', 'add', 'change', 'delete'):
            with self.assertRaises(ValueError):
                register_model_actions(Site, [ModelAction(action_name)])

    def test_empty_action_name_rejected(self):
        with self.assertRaises(ValueError):
            register_model_actions(Site, [ModelAction('')])

    def test_flat_choices_no_duplicates(self):
        """Same action name registered for multiple models should appear once in flat choices."""
        register_model_actions(Device, [ModelAction('render_config')])
        register_model_actions(VirtualMachine, [ModelAction('render_config')])

        device_ct = ObjectType.objects.get_for_model(Device)
        form = ObjectPermissionForm(data={
            'name': 'test',
            'object_types_0': [],
            'object_types_1': [device_ct.pk],
            'can_view': True,
        })
        choices = form.fields['registered_actions'].choices
        names = [name for name, label in choices]
        self.assertEqual(names.count('render_config'), 1)


class ObjectPermissionFormTest(TestCase):

    def setUp(self):
        self._original_actions = dict(registry['model_actions'])
        registry['model_actions'].clear()

    def tearDown(self):
        registry['model_actions'].clear()
        registry['model_actions'].update(self._original_actions)

    def test_shared_action_preselection(self):
        """Editing a permission with render_config should pre-select 'render_config' (plain name)."""
        register_model_actions(Device, [ModelAction('render_config')])
        register_model_actions(VirtualMachine, [ModelAction('render_config')])

        device_ct = ObjectType.objects.get_for_model(Device)
        vm_ct = ObjectType.objects.get_for_model(VirtualMachine)

        permission = ObjectPermission.objects.create(
            name='Test Permission',
            actions=['view', 'render_config'],
        )
        permission.object_types.set([device_ct, vm_ct])

        form = ObjectPermissionForm(instance=permission)

        initial = form.fields['registered_actions'].initial
        self.assertIn('render_config', initial)
        # Should not use the old model-prefixed format
        self.assertNotIn('dcim.device.render_config', initial)
        self.assertNotIn('virtualization.virtualmachine.render_config', initial)
        # Should only appear once despite being registered for two models
        self.assertEqual(initial.count('render_config'), 1)

        # Should not leak into the additional actions field
        self.assertEqual(form.initial['actions'], [])

        permission.delete()

    def test_clean_accepts_valid_registered_action(self):
        """clean() should accept a plain action name when the matching object type is selected."""
        register_model_actions(Device, [ModelAction('render_config')])

        device_ct = ObjectType.objects.get_for_model(Device)
        form = ObjectPermissionForm(data={
            'name': 'test perm',
            'object_types_0': [],
            'object_types_1': [device_ct.pk],
            'registered_actions': ['render_config'],
            'can_view': True,
            'actions': '',
        })
        self.assertTrue(form.is_valid(), form.errors)
        self.assertIn('render_config', form.cleaned_data['actions'])

    def test_clean_rejects_action_without_matching_object_type(self):
        """clean() should reject a registered action when no matching object type is selected."""
        register_model_actions(Device, [ModelAction('render_config')])

        site_ct = ObjectType.objects.get_for_model(Site)
        form = ObjectPermissionForm(data={
            'name': 'test perm',
            'object_types_0': [],
            'object_types_1': [site_ct.pk],
            'registered_actions': ['render_config'],
            'can_view': True,
            'actions': '',
        })
        self.assertFalse(form.is_valid())
        self.assertIn('registered_actions', form.errors)