diff --git a/python/nav/arnold.py b/python/nav/arnold.py index 59a1ce4e85..53f9ab8e6b 100644 --- a/python/nav/arnold.py +++ b/python/nav/arnold.py @@ -31,6 +31,7 @@ from subprocess import Popen, PIPE from collections import namedtuple from smtplib import SMTPException +from typing import Callable from IPy import IP from django.db import connection @@ -40,6 +41,7 @@ from nav.Snmp.errors import AgentError import nav.bitvector import nav.buildconf +from nav.Snmp.profile import get_snmp_session_for_profile from nav.config import find_config_file from nav.errors import GeneralException from nav.models.arnold import Identity, Event @@ -443,7 +445,11 @@ def open_port(identity, username, eventcomment=""): _logger.info("openPort: Port successfully opened") -def change_port_status(action, identity): +def change_port_status( + action, + identity, + agent_getter: Callable = get_snmp_session_for_profile, +): """Use SNMP to change status on an interface. We use ifadminstatus to enable and disable ports @@ -463,14 +469,8 @@ def change_port_status(action, identity): if not profile: raise NoReadWriteManagementProfileError - if not hasattr(profile, "snmp_community") or not hasattr(profile, "snmp_version"): - raise InvalidManagementProfileError - agent = nav.Snmp.Snmp( - host=netbox.ip, - community=profile.snmp_community, - version=profile.snmp_version, - ) + agent = agent_getter(profile)(host=netbox.ip) # Disable or enable based on input try: diff --git a/tests/unittests/arnold/arnold_snmp_test.py b/tests/unittests/arnold/arnold_snmp_test.py index 02d1318478..506234a00f 100644 --- a/tests/unittests/arnold/arnold_snmp_test.py +++ b/tests/unittests/arnold/arnold_snmp_test.py @@ -1,8 +1,10 @@ """Tests for Arnold using snmp objects""" from nav.arnold import change_port_status -from mock import Mock, patch +from unittest.mock import Mock, patch import unittest +from nav.models.manage import ManagementProfile + @patch('nav.Snmp.Snmp', autospec=True) class TestArnoldSnmp(unittest.TestCase): @@ -28,9 +30,13 @@ def create_interface_mock(self): def create_management_profile_mock(self): """Create management profile model mock object""" - profile = Mock() - profile.snmp_version = 1 - profile.snmp_community = "public" + profile = ManagementProfile( + protocol=ManagementProfile.PROTOCOL_SNMP, + configuration={ + "version": 1, + "community": "public", + }, + ) return profile def create_netbox_mock(self): @@ -46,11 +52,9 @@ def test_change_port_status_enable(self, snmp): instance = snmp.return_value identity = Mock() identity.interface = self.interface - change_port_status('enable', identity) + change_port_status('enable', identity, agent_getter=lambda profile: snmp) - snmp.assert_called_once_with( - self.ip, self.profile.snmp_community, self.profile.snmp_version - ) + snmp.assert_called_once_with(self.ip) instance.set.assert_called_once_with( self.port_status_oid + '.' + str(self.ifindex), 'i', 1 ) @@ -60,11 +64,9 @@ def test_change_port_status_disable(self, snmp): instance = snmp.return_value identity = Mock() identity.interface = self.interface - change_port_status('disable', identity) + change_port_status('disable', identity, agent_getter=lambda profile: snmp) - snmp.assert_called_once_with( - self.ip, self.profile.snmp_community, self.profile.snmp_version - ) + snmp.assert_called_once_with(self.ip) instance.set.assert_called_once_with( self.port_status_oid + '.' + str(self.ifindex), 'i', 2 )