Merge pull request #179 from retigra/bug/fix_site_extension

Site extension improvement

netbox_zabbix_sync/modules/cli.py

@@ -133,6 +133,7 @@ def main(arguments):
         zbx_token=zabbix_token,
     )
     syncer.start()
+    syncer.logout()
 
 
 def parse_cli():
@@ -165,7 +166,7 @@ def parse_cli():
         default=None,
     )
     parser.add_argument(
-        "--version", action="version", version="NetBox-Zabbix Sync 3.4.0"
+        "--version", action="version", version="NetBox-Zabbix Sync 4.0.1"
     )
 
     # ── Boolean config overrides ───────────────────────────────────────────────

netbox_zabbix_sync/modules/core.py

@@ -2,6 +2,7 @@
 
 import ssl
 from os import environ
+from pprint import pformat
 from typing import Any
 
 from pynetbox import api as nbapi
@@ -49,6 +50,58 @@ class Sync:
 
         self.config: dict[str, Any] = combined_config
 
+    def _combine_filters(self, config_filter, method_filter):
+        """
+        Combine filters from config and method parameters.
+        Method parameters will overwrite config filters if there are overlaps.
+        """
+        # Check if method filter is provided,
+        # if not return config filter directly
+        combined_filter = config_filter.copy()
+        if method_filter:
+            combined_filter.update(method_filter)
+        return combined_filter
+
+    def _validate_netbox_token(self, token: str, nb_version: str) -> bool:
+        """Validate the format of the NetBox token based on the NetBox version.
+        :param token: The NetBox token to validate.
+        :param nb_version: The version of NetBox being used.
+        :return: True if the token format is valid for the given NetBox version, False otherwise.
+        """
+        support_token_url = (
+            "https://netboxlabs.com/docs/netbox/integrations/rest-api/#v1-and-v2-tokens"  # noqa: S105
+        )
+        token_prefix = "nbt_"  # noqa: S105
+        nb_v2_support_version = "4.5"
+        v2_token = bool(token.startswith(token_prefix) and "." in token)
+        v2_error_token = bool(token.startswith(token_prefix) and "." not in token)
+        # Check if the token is passed without a proper key.token format
+        if v2_error_token:
+            logger.error(
+                "It looks like an invalid v2 token was passed. For more info, see %s",
+                support_token_url,
+            )
+            return False
+        # Warning message for Netbox token v1 with Netbox v4.5 and higher
+        if not v2_token and nb_version >= nb_v2_support_version:
+            logger.warning(
+                "Using Netbox v1 token format. "
+                "Consider updating to a v2 token. For more info, see %s",
+                support_token_url,
+            )
+        elif v2_token and nb_version < nb_v2_support_version:
+            logger.error(
+                "Using Netbox v2 token format with Netbox version lower than 4.5. "
+                "Revert to v1 token or upgrade Netbox to 4.5 or higher. For more info, see %s",
+                support_token_url,
+            )
+            return False
+        elif v2_token and nb_version >= nb_v2_support_version:
+            logger.debug("Using NetBox v2 token format.")
+        else:
+            logger.debug("Using NetBox v1 token format.")
+        return True
+
     def connect(
         self, nb_host, nb_token, zbx_host, zbx_user=None, zbx_pass=None, zbx_token=None
     ):
@@ -117,47 +170,17 @@ class Sync:
             return False
         return True
 
-    def _validate_netbox_token(self, token: str, nb_version: str) -> bool:
-        """Validate the format of the NetBox token based on the NetBox version.
-        :param token: The NetBox token to validate.
-        :param nb_version: The version of NetBox being used.
-        :return: True if the token format is valid for the given NetBox version, False otherwise.
+    def logout(self):
         """
-        support_token_url = (
-            "https://netboxlabs.com/docs/netbox/integrations/rest-api/#v1-and-v2-tokens"  # noqa: S105
-        )
-        token_prefix = "nbt_"  # noqa: S105
-        nb_v2_support_version = "4.5"
-        v2_token = bool(token.startswith(token_prefix) and "." in token)
-        v2_error_token = bool(token.startswith(token_prefix) and "." not in token)
-        # Check if the token is passed without a proper key.token format
-        if v2_error_token:
-            logger.error(
-                "It looks like an invalid v2 token was passed. For more info, see %s",
-                support_token_url,
-            )
-            return False
-        # Warning message for Netbox token v1 with Netbox v4.5 and higher
-        if not v2_token and nb_version >= nb_v2_support_version:
-            logger.warning(
-                "Using Netbox v1 token format. "
-                "Consider updating to a v2 token. For more info, see %s",
-                support_token_url,
-            )
-        elif v2_token and nb_version < nb_v2_support_version:
-            logger.error(
-                "Using Netbox v2 token format with Netbox version lower than 4.5. "
-                "Revert to v1 token or upgrade Netbox to 4.5 or higher. For more info, see %s",
-                support_token_url,
-            )
-            return False
-        elif v2_token and nb_version >= nb_v2_support_version:
-            logger.debug("Using NetBox v2 token format.")
-        else:
-            logger.debug("Using NetBox v1 token format.")
-        return True
+        Logout from Zabbix API
+        """
+        if self.zabbix:
+            self.zabbix.logout()
+            logger.debug("Logged out from Zabbix API.")
+            return True
+        return False
 
-    def start(self):
+    def start(self, device_filter=None, vm_filter=None):
         """
         Run the NetBox to Zabbix sync process.
         """
@@ -196,15 +219,17 @@ class Sync:
         # Set API parameter mapping based on API version
         proxy_name = "host" if str(self.zabbix.version) < "7" else "name"
         # Get all Zabbix and NetBox data
-        netbox_devices = list(
-            self.netbox.dcim.devices.filter(**self.config["nb_device_filter"])
+        dev_filter_combined = self._combine_filters(
+            self.config["nb_device_filter"], device_filter
         )
+        netbox_devices = list(self.netbox.dcim.devices.filter(**dev_filter_combined))
         netbox_vms = []
         if self.config["sync_vms"]:
+            vm_filter_combined = self._combine_filters(
+                self.config["nb_vm_filter"], vm_filter
+            )
             netbox_vms = list(
-                self.netbox.virtualization.virtual_machines.filter(
-                    **self.config["nb_vm_filter"]
-                )
+                self.netbox.virtualization.virtual_machines.filter(**vm_filter_combined)
             )
         netbox_site_groups = convert_recordset(self.netbox.dcim.site_groups.all())
         netbox_regions = convert_recordset(self.netbox.dcim.regions.all())
@@ -258,12 +283,15 @@ class Sync:
                     continue
                 if self.config["extended_site_properties"] and nb_vm.site:
                     logger.debug("Host %s: extending site information.", vm.name)
-                    vm.site = convert_recordset(
-                        self.netbox.dcim.sites.filter(id=nb_vm.site.id)
-                    )
+                    nb_vm.site.full_details()
                 vm.set_inventory(nb_vm)
                 vm.set_usermacros()
                 vm.set_tags()
+                logger.debug(
+                    "Host %s NetBox data: %s",
+                    vm.name,
+                    pformat(dict(nb_vm)),
+                )
                 # Checks if device is in cleanup state
                 if vm.status in self.config["zabbix_device_removal"]:
                     if vm.zabbix_id:
@@ -337,12 +365,14 @@ class Sync:
                     continue
                 if self.config["extended_site_properties"] and nb_device.site:
                     logger.debug("Host %s: extending site information.", device.name)
-                    device.site = convert_recordset(
-                        self.netbox.dcim.sites.filter(id=nb_device.site.id)
-                    )
+                    nb_device.site.full_details()
                 device.set_inventory(nb_device)
                 device.set_usermacros()
                 device.set_tags()
+
+                logger.debug(
+                    "Host %s NetBox data: %s", device.name, pformat(dict(nb_device))
+                )
                 # Checks if device is part of cluster.
                 # Requires clustering variable
                 if device.is_cluster() and self.config["clustering"]:
@@ -401,5 +431,4 @@ class Sync:
                 )
             except SyncError:
                 pass
-        self.zabbix.logout()
         return True

netbox_zabbix_sync/modules/device.py

@@ -62,7 +62,6 @@ class PhysicalDevice:
         self.hostgroups = []
         self.hostgroup_type = "dev"
         self.tenant = nb.tenant
-        self.site = nb.site
         self.config_context = nb.config_context
         self.zbxproxy = None
         self.zabbix_state = 0

tests/test_core.py

@@ -10,7 +10,81 @@ from zabbix_utils import APIRequestError
 from netbox_zabbix_sync.modules.core import Sync
 
 
-class MockNetboxDevice:
+class MockListObject:
+    def __repr__(self):
+        return str(self.object)
+
+    def serialize(self):
+        ret = {k: getattr(self, k) for k in ["object_id", "object_type"]}
+        return ret
+
+    def __getattr__(self, k):
+        return getattr(self.object, k)
+
+    def __iter__(self):
+        for i in ["object_id", "object_type", "object"]:
+            cur_attr = getattr(self, i)
+            if isinstance(cur_attr, MockRecord):
+                cur_attr = dict(cur_attr)
+            yield i, cur_attr
+
+
+class MockRecord:
+    def __init__(self, values, api, endpoint):
+        self.has_details = False
+        self._full_cache = []
+        self._init_cache = []
+        self.api = api
+        self.default_ret = MockRecord
+
+    def __iter__(self):
+        for i in dict(self._init_cache):
+            cur_attr = getattr(self, i)
+            if isinstance(cur_attr, MockRecord):
+                yield i, dict(cur_attr)
+            elif isinstance(cur_attr, list) and all(
+                isinstance(i, (MockRecord, MockListObject)) for i in cur_attr
+            ):
+                yield i, [dict(x) for x in cur_attr]
+            else:
+                yield i, cur_attr
+
+    def __getitem__(self, k):
+        return dict(self)[k]
+
+    def __str__(self):
+        return (
+            getattr(self, "name", None)
+            or getattr(self, "label", None)
+            or getattr(self, "display", None)
+            or ""
+        )
+
+    def __repr__(self):
+        return str(self)
+
+    def __getstate__(self):
+        return self.__dict__
+
+    def __setstate__(self, d):
+        self.__dict__.update(d)
+
+    def __key__(self):
+        if hasattr(self, "id"):
+            return ("mock", self.id)
+        else:
+            return "mock"
+
+    def __hash__(self):
+        return hash(self.__key__())
+
+    def __eq__(self, other):
+        if isinstance(other, MockRecord):
+            return self.__key__() == other.__key__()
+        return NotImplemented
+
+
+class MockNetboxDevice(MockRecord):
     """Mock NetBox device object."""
 
     def __init__(
@@ -31,6 +105,7 @@ class MockNetboxDevice:
         serial="",
         tags=None,
     ):
+        super().__init__(values={}, api=None, endpoint=None)
         self.id = device_id
         self.name = name
         self.status = MagicMock()
@@ -108,7 +183,7 @@ class MockNetboxDevice:
         """Mock save method for NetBox device."""
 
 
-class MockNetboxVM:
+class MockNetboxVM(MockRecord):
     """Mock NetBox virtual machine object.
 
     Mirrors the real NetBox API response structure so the full VirtualMachine
@@ -130,6 +205,7 @@ class MockNetboxVM:
         platform=None,
         tags=None,
     ):
+        super().__init__(values={}, api=None, endpoint=None)
         self.id = vm_id
         self.name = name
         self.status = MagicMock()
@@ -663,49 +739,6 @@ class TestSyncZabbixVersionHandling(unittest.TestCase):
         mock_zabbix.proxygroup.get.assert_not_called()
 
 
-class TestSyncLogout(unittest.TestCase):
-    """Test that sync properly logs out from Zabbix."""
-
-    def _setup_netbox_mock(self, mock_api):
-        """Helper to setup a working NetBox mock."""
-        mock_netbox = MagicMock()
-        mock_api.return_value = mock_netbox
-        mock_netbox.version = "3.5"
-        mock_netbox.extras.custom_fields.filter.return_value = []
-        mock_netbox.dcim.devices.filter.return_value = []
-        mock_netbox.virtualization.virtual_machines.filter.return_value = []
-        mock_netbox.dcim.site_groups.all.return_value = []
-        mock_netbox.dcim.regions.all.return_value = []
-        return mock_netbox
-
-    @patch("netbox_zabbix_sync.modules.core.ZabbixAPI")
-    @patch("netbox_zabbix_sync.modules.core.nbapi")
-    def test_sync_logs_out_from_zabbix(self, mock_api, mock_zabbix_api):
-        """Test that sync calls logout on Zabbix API after completion."""
-        self._setup_netbox_mock(mock_api)
-
-        mock_zabbix = MagicMock()
-        mock_zabbix_api.return_value = mock_zabbix
-        mock_zabbix.version = "6.0"
-        mock_zabbix.hostgroup.get.return_value = []
-        mock_zabbix.template.get.return_value = []
-        mock_zabbix.proxy.get.return_value = []
-
-        syncer = Sync()
-        syncer.connect(
-            "http://netbox.local",
-            "nb_token",
-            "http://zabbix.local",
-            "user",
-            "pass",
-            None,
-        )
-        syncer.start()
-
-        # Verify logout was called
-        mock_zabbix.logout.assert_called_once()
-
-
 class TestSyncProxyNameSanitization(unittest.TestCase):
     """Test proxy name field sanitization for Zabbix 6."""
 
@@ -791,7 +824,6 @@ class TestDeviceHandeling(unittest.TestCase):
         ]
         mock_zabbix.proxy.get.return_value = []
         mock_zabbix.proxygroup.get.return_value = []
-        mock_zabbix.logout = MagicMock()
         # Mock host.get to return empty (host doesn't exist yet)
         mock_zabbix.host.get.return_value = []
         # Mock host.create to return success
@@ -1030,7 +1062,6 @@ class TestDeviceStatusHandling(unittest.TestCase):
         ]
         mock_zabbix.proxy.get.return_value = []
         mock_zabbix.proxygroup.get.return_value = []
-        mock_zabbix.logout = MagicMock()
         mock_zabbix.host.get.return_value = []
         mock_zabbix.host.create.return_value = {"hostids": ["1"]}
         mock_zabbix.host.update.return_value = {"hostids": ["42"]}
@@ -1335,7 +1366,6 @@ class TestVMStatusHandling(unittest.TestCase):
         ]
         mock_zabbix.proxy.get.return_value = []
         mock_zabbix.proxygroup.get.return_value = []
-        mock_zabbix.logout = MagicMock()
         mock_zabbix.host.get.return_value = []
         mock_zabbix.host.create.return_value = {"hostids": ["1"]}
         mock_zabbix.host.update.return_value = {"hostids": ["42"]}
@@ -1581,3 +1611,258 @@ class TestVMStatusHandling(unittest.TestCase):
         syncer.start()
 
         mock_zabbix.host.update.assert_called_once_with(hostid=42, status="1")
+
+
+class TestCombineFilters(unittest.TestCase):
+    """Test the _combine_filters method and filter override behavior in start()."""
+
+    def _setup_netbox_mock(self, mock_api, devices=None, vms=None):
+        """Helper to setup a working NetBox mock."""
+        mock_netbox = MagicMock()
+        mock_api.return_value = mock_netbox
+        mock_netbox.version = "3.5"
+        mock_netbox.extras.custom_fields.filter.return_value = []
+        mock_netbox.dcim.devices.filter.return_value = devices or []
+        mock_netbox.virtualization.virtual_machines.filter.return_value = vms or []
+        mock_netbox.dcim.site_groups.all.return_value = []
+        mock_netbox.dcim.regions.all.return_value = []
+        mock_netbox.extras.journal_entries = MagicMock()
+        return mock_netbox
+
+    def _setup_zabbix_mock(self, mock_zabbix_api, version="7.0"):
+        """Helper to setup a working Zabbix mock."""
+        mock_zabbix = MagicMock()
+        mock_zabbix_api.return_value = mock_zabbix
+        # Set version as float to match expected type in device.py comparisons
+        mock_zabbix.version = float(version)
+        mock_zabbix.hostgroup.get.return_value = [{"groupid": "1", "name": "TestGroup"}]
+        mock_zabbix.template.get.return_value = [
+            {"templateid": "1", "name": "TestTemplate"}
+        ]
+        mock_zabbix.proxy.get.return_value = []
+        mock_zabbix.proxygroup.get.return_value = []
+        mock_zabbix.host.get.return_value = []
+        mock_zabbix.host.create.return_value = {"hostids": ["1"]}
+        return mock_zabbix
+
+    @patch("netbox_zabbix_sync.modules.core.ZabbixAPI")
+    @patch("netbox_zabbix_sync.modules.core.nbapi")
+    def test_filter_override_with_name_parameter(self, mock_api, mock_zabbix_api):
+        """Test that method filter parameter overrides config filter for name.
+
+        Scenario:
+        - Config has nb_device_filter with name="SW01N0"
+        - start() is called with device_filter {"name": "Testdev02"}
+        - Only the device matching "Testdev02" should be processed
+        """
+        # Create two mock devices
+        device_matching_method_filter = MockNetboxDevice(
+            device_id=1, name="Testdev02", status_label="Active"
+        )
+        device_matching_config_filter = MockNetboxDevice(
+            device_id=2, name="SW01N0", status_label="Active"
+        )
+
+        # Setup mocks - the filter should be called with the combined/overridden filter
+        self._setup_netbox_mock(
+            mock_api,
+            devices=[
+                device_matching_method_filter,
+                device_matching_config_filter,
+            ],
+        )
+        self._setup_zabbix_mock(mock_zabbix_api)
+
+        # Create sync with config filter specifying one name
+        syncer = Sync({"nb_device_filter": {"name": "SW01N0"}})
+        syncer.connect(
+            "http://netbox.local",
+            "nb_token",
+            "http://zabbix.local",
+            "user",
+            "pass",
+            None,
+        )
+
+        # Call start with method filter specifying a different name
+        # The method filter should override the config filter
+        syncer.start(device_filter={"name": "Testdev02"})
+
+        # Verify that nbapi.dcim.devices.filter was called with the override filter
+        mock_netbox = mock_api.return_value
+        filter_call_kwargs = mock_netbox.dcim.devices.filter.call_args[1]
+        self.assertEqual(filter_call_kwargs.get("name"), "Testdev02")
+
+    @patch("netbox_zabbix_sync.modules.core.ZabbixAPI")
+    @patch("netbox_zabbix_sync.modules.core.nbapi")
+    def test_filter_override_site_parameter(self, mock_api, mock_zabbix_api):
+        """Test that site filter override works correctly.
+
+        Scenario:
+        - Config has no device filter set
+        - start() is called with device_filter {"site": "fra01"}
+        - Only devices in fra01 site should be processed
+        """
+        # Create mock sites
+        site_fra01 = MagicMock()
+        site_fra01.name = "fra01"
+        site_fra01.slug = "fra01"
+
+        site_ams01 = MagicMock()
+        site_ams01.name = "ams01"
+        site_ams01.slug = "ams01"
+
+        # Create devices in different sites
+        device_fra01 = MockNetboxDevice(
+            device_id=1, name="device-fra01", status_label="Active", site=site_fra01
+        )
+        device_ams01 = MockNetboxDevice(
+            device_id=2, name="device-ams01", status_label="Active", site=site_ams01
+        )
+
+        self._setup_netbox_mock(mock_api, devices=[device_fra01, device_ams01])
+        self._setup_zabbix_mock(mock_zabbix_api)
+
+        syncer = Sync()
+        syncer.connect(
+            "http://netbox.local",
+            "nb_token",
+            "http://zabbix.local",
+            "user",
+            "pass",
+            None,
+        )
+
+        # Call start with site filter for fra01
+        syncer.start(device_filter={"site": "fra01"})
+
+        # Verify that nbapi.dcim.devices.filter was called with the site filter
+        mock_netbox = mock_api.return_value
+        filter_call_kwargs = mock_netbox.dcim.devices.filter.call_args[1]
+        self.assertEqual(filter_call_kwargs.get("site"), "fra01")
+
+    @patch("netbox_zabbix_sync.modules.core.ZabbixAPI")
+    @patch("netbox_zabbix_sync.modules.core.nbapi")
+    def test_config_filter_overridden_by_start_parameter(
+        self, mock_api, mock_zabbix_api
+    ):
+        """Test that start() method filter overrides config filter.
+
+        Scenario:
+        - Config specifies nb_device_filter with {"name": "SW01N0", "site": "ams01"}
+        - start() is called with {"name": "Testdev02"} (only overriding name)
+        - The final filter should be {"name": "Testdev02", "site": "ams01"}
+        - Both name and site filters should be applied with the override
+        """
+        device_matching_all = MockNetboxDevice(
+            device_id=1, name="Testdev02", status_label="Active"
+        )
+
+        self._setup_netbox_mock(mock_api, devices=[device_matching_all])
+        self._setup_zabbix_mock(mock_zabbix_api)
+
+        # Create sync with config filter having multiple parameters
+        syncer = Sync({"nb_device_filter": {"name": "SW01N0", "site": "ams01"}})
+        syncer.connect(
+            "http://netbox.local",
+            "nb_token",
+            "http://zabbix.local",
+            "user",
+            "pass",
+            None,
+        )
+
+        # Call start with method filter that overrides only the name
+        syncer.start(device_filter={"name": "Testdev02"})
+
+        # Verify that nbapi.dcim.devices.filter was called with combined filter
+        # (site from config + name from method parameter)
+        mock_netbox = mock_api.return_value
+        filter_call_kwargs = mock_netbox.dcim.devices.filter.call_args[1]
+        self.assertEqual(filter_call_kwargs.get("name"), "Testdev02")
+        self.assertEqual(filter_call_kwargs.get("site"), "ams01")
+
+    @patch("netbox_zabbix_sync.modules.core.ZabbixAPI")
+    @patch("netbox_zabbix_sync.modules.core.nbapi")
+    def test_vm_filter_override_with_method_parameter(self, mock_api, mock_zabbix_api):
+        """Test that VM filter override works correctly.
+
+        Scenario:
+        - Config enables VM sync with nb_vm_filter {"name": "vm-prod"}
+        - start() is called with vm_filter {"name": "vm-test"}
+        - Only VMs matching "vm-test" should be processed
+        """
+        # Create two mock VMs
+        vm_matching_method_filter = MockNetboxVM(
+            vm_id=1, name="vm-test", status_label="Active"
+        )
+        vm_matching_config_filter = MockNetboxVM(
+            vm_id=2, name="vm-prod", status_label="Active"
+        )
+
+        self._setup_netbox_mock(
+            mock_api,
+            vms=[vm_matching_method_filter, vm_matching_config_filter],
+        )
+        self._setup_zabbix_mock(mock_zabbix_api)
+
+        # Create sync with config filter for VMs
+        syncer = Sync(
+            {
+                "sync_vms": True,
+                "nb_vm_filter": {"name": "vm-prod"},
+            }
+        )
+        syncer.connect(
+            "http://netbox.local",
+            "nb_token",
+            "http://zabbix.local",
+            "user",
+            "pass",
+            None,
+        )
+
+        # Call start with method filter that overrides the VM name filter
+        syncer.start(vm_filter={"name": "vm-test"})
+
+        # Verify that nbapi.virtualization.virtual_machines.filter was called with override
+        mock_netbox = mock_api.return_value
+        filter_call_kwargs = (
+            mock_netbox.virtualization.virtual_machines.filter.call_args[1]
+        )
+        self.assertEqual(filter_call_kwargs.get("name"), "vm-test")
+
+    @patch("netbox_zabbix_sync.modules.core.ZabbixAPI")
+    @patch("netbox_zabbix_sync.modules.core.nbapi")
+    def test_multiple_filter_parameters_combined(self, mock_api, mock_zabbix_api):
+        """Test that multiple filter parameters are correctly combined.
+
+        Scenario:
+        - Config has nb_device_filter with {"site": "fra01", "status": "active"}
+        - start() is called with {"name": "router*"}
+        - The final filter should have all three parameters
+        """
+        device = MockNetboxDevice(device_id=1, name="router01", status_label="Active")
+
+        self._setup_netbox_mock(mock_api, devices=[device])
+        self._setup_zabbix_mock(mock_zabbix_api)
+
+        syncer = Sync({"nb_device_filter": {"site": "fra01", "status": "active"}})
+        syncer.connect(
+            "http://netbox.local",
+            "nb_token",
+            "http://zabbix.local",
+            "user",
+            "pass",
+            None,
+        )
+
+        syncer.start(device_filter={"name": "router*"})
+
+        mock_netbox = mock_api.return_value
+        filter_call_kwargs = mock_netbox.dcim.devices.filter.call_args[1]
+
+        # All three parameters should be present
+        self.assertEqual(filter_call_kwargs.get("site"), "fra01")
+        self.assertEqual(filter_call_kwargs.get("status"), "active")
+        self.assertEqual(filter_call_kwargs.get("name"), "router*")