Support integration ids in all operations

armis_sdk/clients/network_equipment_client.py

@@ -1,4 +1,3 @@
-import asyncio
 from typing import List
 from typing import Set
 
@@ -82,14 +81,13 @@ async def main():
         await self._insert(site.id, new_ids - current_ids)
         await self._delete(site.id, current_ids - new_ids)
 
-    async def _delete(self, site_id: str, network_equipment_device_ids: Set[int]):
+    async def _delete(self, site_id: int, network_equipment_device_ids: Set[int]):
         if not network_equipment_device_ids:
             return
 
         errors = []
         async with self._armis_client.client() as client:
-
-            async def _delete(device_id: int):
+            for device_id in network_equipment_device_ids:
                 response = await client.delete(
                     f"/api/v1/sites/{site_id}/network-equipment/{device_id}/"
                 )
@@ -98,16 +96,14 @@ async def _delete(device_id: int):
                 except HTTPStatusError as error:
                     errors.append(error)
 
-            await asyncio.gather(*map(_delete, network_equipment_device_ids))
-
         if errors:
             raise ResponseError(
                 "Error while deleting network equipment "
                 f"device ids {network_equipment_device_ids!r} from site {site_id!r} ",
                 response_errors=errors,
             )
 
-    async def _insert(self, site_id: str, network_equipment_device_ids: Set[int]):
+    async def _insert(self, site_id: int, network_equipment_device_ids: Set[int]):
         if not network_equipment_device_ids:
             return
 

armis_sdk/clients/site_integrations_client.py

@@ -0,0 +1,105 @@
+from typing import List
+from typing import Set
+
+from httpx import HTTPStatusError
+
+from armis_sdk.core.armis_error import ArmisError
+from armis_sdk.core.armis_error import ResponseError
+from armis_sdk.core.base_entity_client import BaseEntityClient
+from armis_sdk.entities.site import Site
+
+
+class SiteIntegrationsClient(
+    BaseEntityClient
+):  # pylint: disable=too-few-public-methods
+    """
+    A client for interacting with a site's integrations.
+
+    The primary entity for this client is [Site][armis_sdk.entities.site.Site].
+    """
+
+    async def update(self, site: Site):
+        """Update a site's integrations.
+
+        Args:
+            site: The site to update.
+
+        Raises:
+            ResponseError: If an error occurs while communicating with the API.
+            ArmisError: If `site.integration_ids` is not set.
+
+        Example:
+            ```python linenums="1" hl_lines="9"
+            import asyncio
+
+            from armis_sdk.clients.site_integrations_client import SiteIntegrationsClient
+
+
+            async def main():
+                site_integrations_client = SiteIntegrationsClient()
+                site = Site(id=1, integration_ids=[1, 2, 3])
+                await site_integrations_client.update(site)
+
+            asyncio.run(main())
+            ```
+        """
+
+        if site.integration_ids is None:
+            raise ArmisError("The property 'integration_ids' must be set.")
+
+        new_ids = set(site.integration_ids)
+        current_ids = set(await self._list(site.id))
+
+        await self._insert(site.id, new_ids - current_ids)
+        await self._delete(site.id, current_ids - new_ids)
+
+    async def _delete(self, site_id: int, integration_ids: Set[int]):
+        if not integration_ids:
+            return
+
+        errors = []
+        async with self._armis_client.client() as client:
+            for integration_id in integration_ids:
+                response = await client.delete(
+                    f"/api/v1/sites/{site_id}/integrations-ids/{integration_id}/"
+                )
+                try:
+                    response.raise_for_status()
+                except HTTPStatusError as error:
+                    errors.append(error)
+
+        if errors:
+            raise ResponseError(
+                "Error while deleting integration ids "
+                f"{integration_ids!r} from site {site_id!r} ",
+                response_errors=errors,
+            )
+
+    async def _insert(self, site_id: int, integration_ids: Set[int]):
+        if not integration_ids:
+            return
+
+        errors = []
+        async with self._armis_client.client() as client:
+            for integration_id in integration_ids:
+                response = await client.post(
+                    f"/api/v1/sites/{site_id}/integrations-ids/",
+                    json={"integrationId": integration_id},
+                )
+                try:
+                    response.raise_for_status()
+                except HTTPStatusError as error:
+                    errors.append(error)
+
+        if errors:
+            raise ResponseError(
+                "Error while inserting integration ids "
+                f"{integration_ids!r} to site {site_id!r} ",
+                response_errors=errors,
+            )
+
+    async def _list(self, site_id: int) -> List[int]:
+        async with self._armis_client.client() as client:
+            response = await client.get(f"/api/v1/sites/{site_id}/integrations-ids/")
+            data = self._get_data(response)
+        return data["integrationIds"]

armis_sdk/clients/sites_client.py

@@ -2,6 +2,7 @@
 from typing import List
 
 from armis_sdk.clients.network_equipment_client import NetworkEquipmentClient
+from armis_sdk.clients.site_integrations_client import SiteIntegrationsClient
 from armis_sdk.core import response_utils
 from armis_sdk.core.armis_error import ArmisError
 from armis_sdk.core.base_entity_client import BaseEntityClient
@@ -17,11 +18,13 @@ class SitesClient(BaseEntityClient):
 
     Attributes:
         network_equipment_client (NetworkEquipmentClient): An instance of [NetworkEquipmentClient][armis_sdk.clients.network_equipment_client.NetworkEquipmentClient]
+        site_integrations_client (SiteIntegrationsClient): An instance of [SiteIntegrationsClient][armis_sdk.clients.site_integrations_client.SiteIntegrationsClient]
     """
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.network_equipment_client = NetworkEquipmentClient(self._armis_client)
+        self.site_integrations_client = SiteIntegrationsClient(self._armis_client)
 
     async def create(self, site: Site) -> Site:
         """Create a `Site`.
@@ -247,7 +250,12 @@ async def main():
 
         data = site.model_dump(
             by_alias=True,
-            exclude={"children", "id", "network_equipment_device_ids"},
+            exclude={
+                "children",
+                "id",
+                "integration_ids",
+                "network_equipment_device_ids",
+            },
             exclude_none=True,
         )
 
@@ -258,3 +266,6 @@ async def main():
 
         if site.network_equipment_device_ids is not None:
             await self.network_equipment_client.update(site)
+
+        if site.integration_ids is not None:
+            await self.site_integrations_client.update(site)

armis_sdk/entities/site.py

@@ -53,6 +53,11 @@ class Site(BaseEntity):
     ] = None
     """The ids of network equipment devices associated with the site."""
 
+    integration_ids: Annotated[
+        Optional[List[int]], BeforeValidator(ensure_list_of_ints)
+    ] = None
+    """The ids of the integration associated with the site."""
+
     children: Optional[List["Site"]] = Field(default_factory=list)
     """The sub-sites that are directly under this site 
     (their `parent_id` will match this site's `id`)."""

docs/clients/SiteIntegrationsClient.md

@@ -0,0 +1 @@
+::: armis_sdk.clients.site_integrations_client.SiteIntegrationsClient

mkdocs.yml

@@ -8,7 +8,6 @@ nav:
       - AsqRule: entities/AsqRule.md
       - Site: entities/Site.md
   - Clients:
-      - NetworkEquipmentClient: clients/NetworkEquipmentClient.md
       - SitesClient: clients/SitesClient.md
   - Core:
       - ArmisClient: core/ArmisClient.md

pylintrc

@@ -5,6 +5,7 @@ ignore-paths=
     .venv
 
 disable=
+    duplicate-code,
     missing-class-docstring,
     missing-module-docstring,
     missing-function-docstring,

tests/armis_sdk/clients/network_equipment_client_test.py

@@ -50,14 +50,14 @@ async def test_update(httpx_mock: pytest_httpx.HTTPXMock):
 
     network_equipment_client = NetworkEquipmentClient()
 
-    site = Site(id="1", name="mock_site", network_equipment_device_ids=[2, 4, 5, 6])
+    site = Site(id=1, name="mock_site", network_equipment_device_ids=[2, 4, 5, 6])
 
     await network_equipment_client.update(site)
 
 
 @pytest.mark.usefixtures("setup_env_variables")
 async def test_update_with_no_devices_set():
-    site = Site(id="1", name="mock_site")
+    site = Site(id=1, name="mock_site")
     network_equipment_client = NetworkEquipmentClient()
 
     with pytest.raises(

tests/armis_sdk/clients/site_integrations_client_test.py

@@ -0,0 +1,55 @@
+import pytest
+import pytest_httpx
+
+from armis_sdk.clients.site_integrations_client import SiteIntegrationsClient
+from armis_sdk.core.armis_error import ArmisError
+from armis_sdk.entities.site import Site
+
+pytest_plugins = ["tests.plugins.setup_plugin"]
+
+
+@pytest.mark.usefixtures("setup_env_variables", "authorized")
+async def test_update(httpx_mock: pytest_httpx.HTTPXMock):
+    # List current ids
+    httpx_mock.add_response(
+        url="https://mock_tenant.armis.com/api/v1/sites/1/integrations-ids/",
+        method="GET",
+        json={"data": {"integrationIds": [1, 2, 3, 4]}},
+    )
+
+    # Add new ids
+    httpx_mock.add_response(
+        url="https://mock_tenant.armis.com/api/v1/sites/1/integrations-ids/",
+        method="POST",
+        match_json={"integrationId": 5},
+    )
+    httpx_mock.add_response(
+        url="https://mock_tenant.armis.com/api/v1/sites/1/integrations-ids/",
+        method="POST",
+        match_json={"integrationId": 6},
+    )
+
+    # Remove old ids
+    httpx_mock.add_response(
+        url="https://mock_tenant.armis.com/api/v1/sites/1/integrations-ids/1/",
+        method="DELETE",
+    )
+    httpx_mock.add_response(
+        url="https://mock_tenant.armis.com/api/v1/sites/1/integrations-ids/3/",
+        method="DELETE",
+    )
+
+    site_integrations_client = SiteIntegrationsClient()
+
+    site = Site(id=1, name="mock_site", integration_ids=[2, 4, 5, 6])
+
+    await site_integrations_client.update(site)
+
+
+@pytest.mark.usefixtures("setup_env_variables")
+async def test_update_with_no_integrations_set():
+    site = Site(id=1, name="mock_site")
+    site_integrations_client = SiteIntegrationsClient()
+
+    with pytest.raises(ArmisError, match="The property 'integration_ids' must be set."):
+        await site_integrations_client.update(site)

tests/armis_sdk/clients/sites_client_test.py

@@ -13,7 +13,12 @@ async def test_create(httpx_mock: pytest_httpx.HTTPXMock):
     httpx_mock.add_response(
         url="https://mock_tenant.armis.com/api/v1/sites/",
         method="POST",
-        match_json={"name": "mock_site", "location": "mock_location", "parentId": 2},
+        match_json={
+            "name": "mock_site",
+            "location": "mock_location",
+            "parentId": 2,
+            "integrationIds": [4, 5, 6],
+        },
         json={"data": {"id": "1"}},
     )
     httpx_mock.add_response(
@@ -27,6 +32,7 @@ async def test_create(httpx_mock: pytest_httpx.HTTPXMock):
         location="mock_location",
         parent_id=2,
         network_equipment_device_ids=[1, 2, 3],
+        integration_ids=[4, 5, 6],
     )
 
     sites_client = SitesClient()
@@ -38,6 +44,7 @@ async def test_create(httpx_mock: pytest_httpx.HTTPXMock):
         location="mock_location",
         parent_id=2,
         network_equipment_device_ids=[1, 2, 3],
+        integration_ids=[4, 5, 6],
     )
 
 
@@ -95,6 +102,8 @@ async def test_get(httpx_mock: pytest_httpx.HTTPXMock):
                 "id": "1",
                 "name": "mock_site_1",
                 "ruleAql": '{"or": ["asq1", "asq2"]}',
+                "networkEquipmentDeviceIds": ["1", "2", "3"],
+                "integrationIds": ["4", "5", "6"],
             }
         },
     )
@@ -103,7 +112,11 @@ async def test_get(httpx_mock: pytest_httpx.HTTPXMock):
     site = await sites_client.get("1")
 
     assert site == Site(
-        id=1, name="mock_site_1", asq_rule=AsqRule(or_=["asq1", "asq2"])
+        id=1,
+        name="mock_site_1",
+        asq_rule=AsqRule(or_=["asq1", "asq2"]),
+        network_equipment_device_ids=[1, 2, 3],
+        integration_ids=[4, 5, 6],
     )
 
 
@@ -192,6 +205,7 @@ async def test_hierarchy(httpx_mock: pytest_httpx.HTTPXMock):
                 "location": "mock_location",
                 "parentId": "1",
                 "tier": "mock_tier",
+                "integrationIds": ["4", "5", "6"],
                 "networkEquipmentDeviceIds": ["7", "8", "9"],
             },
             Site(
@@ -202,6 +216,7 @@ async def test_hierarchy(httpx_mock: pytest_httpx.HTTPXMock):
                 location="mock_location",
                 parent_id=1,
                 tier="mock_tier",
+                integration_ids=[4, 5, 6],
                 network_equipment_device_ids=[7, 8, 9],
             ),
             id="All fields",
@@ -320,6 +335,37 @@ async def test_update_with_network_equipment_device_ids(
     await sites_client.update(site)
 
 
+async def test_update_with_integration_id(httpx_mock: pytest_httpx.HTTPXMock):
+    # List current ids
+    httpx_mock.add_response(
+        url="https://mock_tenant.armis.com/api/v1/sites/1/integrations-ids/",
+        method="GET",
+        json={"data": {"integrationIds": []}},
+    )
+
+    # Add new ids
+    httpx_mock.add_response(
+        url="https://mock_tenant.armis.com/api/v1/sites/1/integrations-ids/",
+        method="POST",
+        match_json={"integrationId": 1},
+    )
+    httpx_mock.add_response(
+        url="https://mock_tenant.armis.com/api/v1/sites/1/integrations-ids/",
+        method="POST",
+        match_json={"integrationId": 2},
+    )
+    httpx_mock.add_response(
+        url="https://mock_tenant.armis.com/api/v1/sites/1/integrations-ids/",
+        method="POST",
+        match_json={"integrationId": 3},
+    )
+
+    sites_client = SitesClient()
+    site = Site(id=1, integration_ids=[1, 2, 3])
+
+    await sites_client.update(site)
+
+
 async def test_update_without_id(httpx_mock: pytest_httpx.HTTPXMock):
     httpx_mock.reset()