diff --git a/pynetbox_data_uploader/lib/netbox_api/netbox_get_id.py b/pynetbox_data_uploader/lib/netbox_api/netbox_get_id.py index 8d30f514..65162764 100644 --- a/pynetbox_data_uploader/lib/netbox_api/netbox_get_id.py +++ b/pynetbox_data_uploader/lib/netbox_api/netbox_get_id.py @@ -1,7 +1,6 @@ from operator import attrgetter from typing import Union, Dict from lib.enums.dcim_device_id import DeviceInfoID -from lib.enums.dcim_device_no_id import DeviceInfoNoID # pylint:disable = too-few-public-methods @@ -13,11 +12,9 @@ class NetboxGetID: def __init__(self, api): """ - This method allows the Netbox Api Object and Enums to be accessible within the class. + This method allows the Netbox Api Object to be accessible within the class. """ self.netbox = api - self.enums_id = DeviceInfoID - self.enums_no_id = DeviceInfoNoID def get_id( self, attr_string: str, netbox_value: str, site_value: str @@ -30,7 +27,7 @@ def get_id( :return: Returns the value/ID """ attr_string = attr_string.upper() - attr_to_look_for = getattr(self.enums_id, attr_string).value # Gets enums value + attr_to_look_for = DeviceInfoID[attr_string].value # Gets enums value value = attrgetter(attr_to_look_for)(self.netbox) # Gets netbox attr if attr_string == "DEVICE_TYPE": value = value.get(slug=netbox_value).id @@ -53,7 +50,7 @@ def get_id_from_key(self, key: str, dictionary: Dict) -> Union[str, int]: :param dictionary: The device dictionary being referenced. :return: If an ID was needed and found it returns the ID. If an ID was not needed it returns the original value. """ - if key.upper() not in list(self.enums_no_id.__members__): + if key.upper() in [prop.name for prop in DeviceInfoID]: value = self.get_id( attr_string=key, netbox_value=dictionary[key], diff --git a/pynetbox_data_uploader/tests/test_netbox_get_id.py b/pynetbox_data_uploader/tests/test_netbox_get_id.py index 4b7287b1..424debae 100644 --- a/pynetbox_data_uploader/tests/test_netbox_get_id.py +++ b/pynetbox_data_uploader/tests/test_netbox_get_id.py @@ -1,6 +1,8 @@ -from unittest.mock import NonCallableMock +from unittest.mock import NonCallableMock, patch import pytest from lib.netbox_api.netbox_get_id import NetboxGetID +from lib.enums.dcim_device_id import DeviceInfoID +from lib.enums.dcim_device_no_id import DeviceInfoNoID @pytest.fixture(name="instance") @@ -11,3 +13,29 @@ def instance_fixture(): """ netbox = NonCallableMock() return NetboxGetID(netbox) + + +def test_get_id_from_key_with_id_enums(instance): + """ + This test ensures that the get_id method is called for all properties in the DeviceInfoID enum. + """ + with patch("lib.netbox_api.netbox_get_id.NetboxGetID.get_id") as mock_get_id: + for member in [prop.name for prop in DeviceInfoID]: + mock_dictionary = {member: "abc", + "site": "def"} + res = instance.get_id_from_key(key=member, dictionary=mock_dictionary) + mock_get_id.assert_called() + assert res == mock_get_id.return_value + + +def test_get_id_from_key_with_no_id_enums(instance): + """ + This test ensures that the get_id method is not called for all properties in the DeviceInfoNoID enum. + """ + with patch("lib.netbox_api.netbox_get_id.NetboxGetID.get_id") as mock_get_id: + for member in [prop.name for prop in DeviceInfoNoID]: + mock_dictionary = {member: "abc", + "site": "def"} + res = instance.get_id_from_key(key=member, dictionary=mock_dictionary) + mock_get_id.assert_not_called() + assert res == mock_dictionary[member]