Improved NetboxClient.

gso/services/netbox_client.py

@@ -3,7 +3,7 @@ from uuid import UUID
 
 import pydantic
 import pynetbox
-from pynetbox.models.dcim import Devices, DeviceTypes, InterfaceConnection, Interfaces
+from pynetbox.models.dcim import Devices, DeviceTypes, Interfaces
 
 from gso.products.product_types.router import Router
 from gso.settings import load_oss_params
@@ -62,8 +62,19 @@ class NetboxClient:
         )
 
     def get_device_by_name(self, device_name: str) -> Devices:
-        """Return the device object by name from netbox, or ``None`` if not found."""
-        return self.netbox.dcim.devices.get(name=device_name)
+        """Return the device object by name from netbox, or raise not found."""
+        device = self.netbox.dcim.devices.get(name=device_name)
+        if device is None:
+            raise NotFoundError(f"Device: {device_name} not found.")
+        return device
+
+    def get_interface_by_name_and_device(self, iface_name: str, device_name: str) -> Interfaces:
+        """Return the interface lists by name and device name from netbox."""
+        device = self.get_device_by_name(device_name)
+        interface = self.netbox.dcim.interfaces.get(device_id=device.id, name=iface_name)
+        if interface is None:
+            raise NotFoundError(f"Interface: {iface_name} on device with id: {device.id} not found.")
+        return interface
 
     def get_interfaces_by_device(self, device_name: str, speed: str) -> list[Interfaces]:
         """Get all interfaces of a device by name and speed that are not reserved and not allocated."""
@@ -95,8 +106,7 @@ class NetboxClient:
     def delete_interface(self, device_name: str, iface_name: str) -> None:
         """Delete an interface from a device by name."""
 
-        device = self.get_device_by_name(device_name)
-        interface = self.netbox.dcim.interfaces.get(device_id=device.id, name=iface_name)
+        interface = self.get_interface_by_name_and_device(iface_name, device_name)
         return interface.delete()
 
     def create_device_type(self, manufacturer: str, model: str, slug: str) -> DeviceTypes:
@@ -179,14 +189,10 @@ class NetboxClient:
 
         Returns the interface object after assignment.
         """
-        # Get device id
-        device = self.get_device_by_name(device_name)
-
-        # Get interface for device
-        iface = self.netbox.dcim.interfaces.get(name=iface_name, device_id=device.id)
+        iface = self.get_interface_by_name_and_device(iface_name, device_name)
 
         # Get LAG
-        lag = self.netbox.dcim.interfaces.get(name=lag_name, device_id=device.id)
+        lag = self.get_interface_by_name_and_device(lag_name, device_name)
 
         # Assign interface to LAG, ensuring it doesn't already belong to a LAG
         if iface.lag:
@@ -206,12 +212,7 @@ class NetboxClient:
         """Reserve an interface by enabling it."""
 
         # First get interface from device
-        device = self.get_device_by_name(device_name)
-        interface = self.netbox.dcim.interfaces.get(device_id=device.id, name=iface_name)
-
-        # Check if interface exists
-        if interface is None:
-            raise NotFoundError(f"Interface: {iface_name} on device: {device_name} not found.")
+        interface = self.get_interface_by_name_and_device(iface_name, device_name)
 
         # Check if interface is reserved
         if interface.enabled:
@@ -227,12 +228,7 @@ class NetboxClient:
         """Allocate an interface by marking it as connected."""
 
         # First get interface from device
-        device = self.get_device_by_name(device_name)
-        interface = self.netbox.dcim.interfaces.get(device_id=device.id, name=iface_name)
-
-        # Check if interface exists
-        if interface is None:
-            raise NotFoundError(f"Interface: {iface_name} on device: {device_name} not found.")
+        interface = self.get_interface_by_name_and_device(iface_name, device_name)
 
         # Check if interface is reserved
         if interface.mark_connected:
@@ -248,13 +244,7 @@ class NetboxClient:
         """Free interface by marking disconnect and disable it."""
 
         # First get interface from device
-        device = self.get_device_by_name(device_name)
-        interface = self.netbox.dcim.interfaces.get(device_id=device.id, name=iface_name)
-
-        # Check if interface is available
-        if interface is None:
-            raise NotFoundError(f"Interface: {iface_name} on device: {device_name} not found.")
-
+        interface = self.get_interface_by_name_and_device(iface_name, device_name)
         interface.mark_connected = False
         interface.enabled = False
         interface.description = ""
@@ -307,13 +297,3 @@ class NetboxClient:
         return self.netbox.dcim.interfaces.filter(
             device=device.name, enabled=False, mark_connected=False, speed=speed_bps
         )
-
-    def get_interface_by_name_and_device(self, router_id: UUID, interface_name: str) -> InterfaceConnection:
-        """Return the interface object by name and device from netbox, or ``None`` if not found."""
-
-        router = Router.from_subscription(router_id).router.router_fqdn
-        device = self.get_device_by_name(router)
-        try:
-            return self.netbox.dcim.interfaces.get(device=device.name, name=interface_name)
-        except pynetbox.RequestError:
-            raise NotFoundError(f"Interface: {interface_name} on device: {device.name} not found.")

gso/utils/helpers.py

@@ -65,7 +65,9 @@ def available_interfaces_choices_including_current_members(
     available_interfaces = list(NetboxClient().get_available_interfaces(router_id, speed))
     available_interfaces.extend(
         [
-            NetboxClient().get_interface_by_name_and_device(router_id, interface.interface_name)
+            NetboxClient().get_interface_by_name_and_device(
+                interface.interface_name, Router.from_subscription(router_id).router.router_fqdn
+            )
             for interface in interfaces
         ]
     )

test/workflows/iptrunk/test_migrate_iptrunk.py

@@ -18,7 +18,6 @@ from test.workflows.iptrunk.test_create_iptrunk import MockedNetboxClient
 @pytest.mark.workflow
 @patch("gso.workflows.iptrunk.migrate_iptrunk.provisioning_proxy.migrate_ip_trunk")
 @patch("gso.workflows.iptrunk.migrate_iptrunk.provisioning_proxy.provision_ip_trunk")
-@patch("gso.services.netbox_client.NetboxClient.get_device_by_name")
 @patch("gso.services.netbox_client.NetboxClient.get_available_interfaces")
 @patch("gso.services.netbox_client.NetboxClient.get_available_lags")
 @patch("gso.services.netbox_client.NetboxClient.create_interface")
@@ -36,7 +35,6 @@ def test_migrate_iptrunk_success(
     mocked_create_interface,
     mocked_get_available_lags,
     mocked_get_available_interfaces,
-    mocked_get_device_by_name,
     mock_provision_ip_trunk,
     mock_migrate_ip_trunk,
     iptrunk_subscription_factory,
@@ -45,7 +43,6 @@ def test_migrate_iptrunk_success(
 ):
     #  Set up mock return values
     mocked_netbox = MockedNetboxClient()
-    mocked_get_device_by_name.return_value = mocked_netbox.get_device_by_name()
     mocked_get_available_interfaces.return_value = mocked_netbox.get_available_interfaces()
     mocked_attach_interface_to_lag.return_value = mocked_netbox.attach_interface_to_lag()
     mocked_reserve_interface.return_value = mocked_netbox.reserve_interface()

test/workflows/iptrunk/test_modify_trunk_interface.py

@@ -17,7 +17,6 @@ from test.workflows.iptrunk.test_create_iptrunk import MockedNetboxClient
 
 @pytest.mark.workflow
 @patch("gso.workflows.iptrunk.modify_trunk_interface.provisioning_proxy.provision_ip_trunk")
-@patch("gso.services.netbox_client.NetboxClient.get_device_by_name")
 @patch("gso.services.netbox_client.NetboxClient.get_available_interfaces")
 @patch("gso.services.netbox_client.NetboxClient.attach_interface_to_lag")
 @patch("gso.services.netbox_client.NetboxClient.reserve_interface")
@@ -31,14 +30,12 @@ def test_iptrunk_modify_trunk_interface_success(
     mocked_reserve_interface,
     mocked_attach_interface_to_lag,
     mocked_get_available_interfaces,
-    mocked_get_device_by_name,
     mock_provision_ip_trunk,
     iptrunk_subscription_factory,
     faker,
 ):
     #  Set up mock return values
     mocked_netbox = MockedNetboxClient()
-    mocked_get_device_by_name.return_value = mocked_netbox.get_device_by_name()
     mocked_get_available_interfaces.return_value = mocked_netbox.get_available_interfaces()
     mocked_attach_interface_to_lag.return_value = mocked_netbox.attach_interface_to_lag()
     mocked_reserve_interface.return_value = mocked_netbox.reserve_interface()

test/workflows/iptrunk/test_terminate_iptrunk.py

@@ -18,13 +18,11 @@ from test.workflows import (
 @patch("gso.workflows.iptrunk.terminate_iptrunk.provisioning_proxy.provision_ip_trunk")
 @patch("gso.workflows.iptrunk.terminate_iptrunk.provisioning_proxy.deprovision_ip_trunk")
 @patch("gso.workflows.iptrunk.terminate_iptrunk.infoblox.delete_network")
-@patch("gso.services.netbox_client.NetboxClient.get_device_by_name")
 @patch("gso.services.netbox_client.NetboxClient.delete_interface")
 @patch("gso.services.netbox_client.NetboxClient.free_interface")
 def test_successful_iptrunk_termination(
     mocked_free_interface,
     mocked_delete_interface,
-    mocked_get_device_by_name,
     mock_infoblox_delete_network,
     mock_deprovision_ip_trunk,
     mock_provision_ip_trunk,
@@ -34,7 +32,6 @@ def test_successful_iptrunk_termination(
     #  Set up mock return values
     product_id = iptrunk_subscription_factory()
     mocked_netbox = MockedNetboxClient()
-    mocked_get_device_by_name.return_value = mocked_netbox.get_device_by_name()
     mocked_delete_interface.return_value = mocked_netbox.delete_interface()
     mocked_free_interface.return_value = mocked_netbox.free_interface()