From 5a5923dcda14adbc8ee0aafdccaef4ea49925ea4 Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Mon, 27 Apr 2026 08:05:06 +0000 Subject: [PATCH] Auto commit from main repo: manual-sync --- armis_sdk/clients/assets_client.py | 2 +- armis_sdk/core/armis_client.py | 19 +++++-- pyproject.toml | 2 +- tests/armis_sdk/clients/assets_client_test.py | 53 ++++++++++++------- 4 files changed, 51 insertions(+), 25 deletions(-) diff --git a/armis_sdk/clients/assets_client.py b/armis_sdk/clients/assets_client.py index f9a0d92..20ac34f 100644 --- a/armis_sdk/clients/assets_client.py +++ b/armis_sdk/clients/assets_client.py @@ -501,7 +501,7 @@ async def _list_assets( "fields": fields, "filter": filter_, } - async for item in self._armis_client.list("/v3/assets/_search", body=body, after_location="filter"): + async for item in self._armis_client.list("/v3/assets/_search", body=body, pagination_location="filter"): yield asset_class.from_search_result(item) @classmethod diff --git a/armis_sdk/core/armis_client.py b/armis_sdk/core/armis_client.py index 9d1d4de..85e6571 100644 --- a/armis_sdk/core/armis_client.py +++ b/armis_sdk/core/armis_client.py @@ -89,13 +89,18 @@ def client(self, retries: int | None = None, backoff: float | None = None): trust_env=True, ) - async def list(self, url: str, body: dict | None = None, after_location: str | None = None) -> AsyncIterator[dict]: + async def list( + self, + url: str, + body: dict | None = None, + pagination_location: str | None = None, + ) -> AsyncIterator[dict]: """List all items from a paginated endpoint. Args: url (str): The relative endpoint URL. body (dict): Payload to send as POST request. - after_location (str): The nested object location to use for pagination. + pagination_location (str): The nested object location to use for pagination. Returns: An (async) iterator of `dict`s. @@ -123,7 +128,11 @@ async def main(): """ page_size = int(os.getenv(ARMIS_PAGE_SIZE, str(DEFAULT_PAGE_LENGTH))) async with self.client() as client: - params = {"limit": page_size, **(body or {})} + params = {**(body or {})} + if pagination_location: + params[pagination_location]["limit"] = page_size + else: + params["limit"] = page_size while True: if body: response = await client.post(url, json=params) @@ -134,8 +143,8 @@ async def main(): for item in items: yield item if next_ := data.get("next"): - if after_location: - params[after_location]["after"] = next_ + if pagination_location: + params[pagination_location]["after"] = next_ else: params["after"] = next_ else: diff --git a/pyproject.toml b/pyproject.toml index 5f48443..d21faf0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "armis_sdk" -version = "1.2.1" +version = "1.2.2" description = "The Armis SDK is a package that encapsulates common use-cases for interacting with the Armis platform." authors = [ { name = "Shai Lachmanovich", email = "shai@armis.com" }, diff --git a/tests/armis_sdk/clients/assets_client_test.py b/tests/armis_sdk/clients/assets_client_test.py index 640af4a..2214912 100644 --- a/tests/armis_sdk/clients/assets_client_test.py +++ b/tests/armis_sdk/clients/assets_client_test.py @@ -25,12 +25,12 @@ async def test_list_by_last_seen_datetime(httpx_mock: pytest_httpx.HTTPXMock): url="https://api.armis.com/v3/assets/_search", method="POST", match_json={ - "limit": 100, "asset_type": "DEVICE", "fields": assets_test_data.ALL_DEVICE_FIELDS, "filter": { "filter_criteria": "LAST_SEEN", "last_seen_ge": "2025-12-03T00:00:00", + "limit": 100, }, }, json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]}, @@ -50,12 +50,12 @@ async def test_list_by_last_seen_datetime_explicit_fields( url="https://api.armis.com/v3/assets/_search", method="POST", match_json={ - "limit": 100, "asset_type": "DEVICE", "fields": ["brand", "custom.MyField1", "custom.MyField2", "purdue_level"], "filter": { "filter_criteria": "LAST_SEEN", "last_seen_ge": "2025-12-03T00:00:00", + "limit": 100, }, }, json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_PARTIAL_RAW_DATA}]}, @@ -74,10 +74,13 @@ async def test_list_by_last_seen_timedelta(httpx_mock: pytest_httpx.HTTPXMock): url="https://api.armis.com/v3/assets/_search", method="POST", match_json={ - "limit": 100, "asset_type": "DEVICE", "fields": assets_test_data.ALL_DEVICE_FIELDS, - "filter": {"filter_criteria": "LAST_SEEN", "last_seen_seconds": 3600}, + "filter": { + "filter_criteria": "LAST_SEEN", + "last_seen_seconds": 3600, + "limit": 100, + }, }, json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]}, ) @@ -96,10 +99,13 @@ async def test_list_by_last_seen_timedelta_explicit_fields( url="https://api.armis.com/v3/assets/_search", method="POST", match_json={ - "limit": 100, "asset_type": "DEVICE", "fields": ["brand", "custom.MyField1", "custom.MyField2", "purdue_level"], - "filter": {"filter_criteria": "LAST_SEEN", "last_seen_seconds": 3600}, + "filter": { + "filter_criteria": "LAST_SEEN", + "last_seen_seconds": 3600, + "limit": 100, + }, }, json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_PARTIAL_RAW_DATA}]}, ) @@ -130,13 +136,13 @@ async def test_list_by_asset_id(httpx_mock: pytest_httpx.HTTPXMock): url="https://api.armis.com/v3/assets/_search", method="POST", match_json={ - "limit": 100, "asset_type": "DEVICE", "fields": assets_test_data.ALL_DEVICE_FIELDS, "filter": { "filter_criteria": "ASSET_ID", "asset_id_source": "IPV4_ADDRESS", "asset_ids": ["1.1.1.1"], + "limit": 100, }, }, json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]}, @@ -161,13 +167,13 @@ async def test_list_by_asset_id_explicit_fields(httpx_mock: pytest_httpx.HTTPXMo url="https://api.armis.com/v3/assets/_search", method="POST", match_json={ - "limit": 100, "asset_type": "DEVICE", "fields": ["brand", "custom.MyField1", "custom.MyField2", "purdue_level"], "filter": { "filter_criteria": "ASSET_ID", "asset_id_source": "IPV4_ADDRESS", "asset_ids": ["1.1.1.1"], + "limit": 100, }, }, json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_PARTIAL_RAW_DATA}]}, @@ -205,10 +211,13 @@ async def test_list_by_boundary_id(httpx_mock: pytest_httpx.HTTPXMock): url="https://api.armis.com/v3/assets/_search", method="POST", match_json={ - "limit": 100, "asset_type": "DEVICE", "fields": assets_test_data.ALL_DEVICE_FIELDS, - "filter": {"filter_criteria": "BOUNDARY_ID", "boundary_ids": [1, 2, 3]}, + "filter": { + "filter_criteria": "BOUNDARY_ID", + "boundary_ids": [1, 2, 3], + "limit": 100, + }, }, json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]}, ) @@ -253,10 +262,13 @@ async def test_list_by_multiple(httpx_mock: pytest_httpx.HTTPXMock, kwargs, expe url="https://api.armis.com/v3/assets/_search", method="POST", match_json={ - "limit": 100, "asset_type": "DEVICE", "fields": assets_test_data.ALL_DEVICE_FIELDS, - "filter": {"filter_criteria": "MULTIPLE", "filters": expected_filters}, + "filter": { + "filter_criteria": "MULTIPLE", + "filters": expected_filters, + "limit": 100, + }, }, json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]}, ) @@ -313,10 +325,13 @@ async def test_list_by_site_id(httpx_mock: pytest_httpx.HTTPXMock): url="https://api.armis.com/v3/assets/_search", method="POST", match_json={ - "limit": 100, "asset_type": "DEVICE", "fields": assets_test_data.ALL_DEVICE_FIELDS, - "filter": {"filter_criteria": "SITE_ID", "site_ids": [1, 2, 3]}, + "filter": { + "filter_criteria": "SITE_ID", + "site_ids": [1, 2, 3], + "limit": 100, + }, }, json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]}, ) @@ -499,10 +514,13 @@ async def test_list_assets_pagination(monkeypatch, httpx_mock: pytest_httpx.HTTP url="https://api.armis.com/v3/assets/_search", method="POST", match_json={ - "limit": 1, "asset_type": "DEVICE", "fields": assets_test_data.ALL_DEVICE_FIELDS, - "filter": {"filter_criteria": "LAST_SEEN", "last_seen_seconds": 3600}, + "filter": { + "filter_criteria": "LAST_SEEN", + "last_seen_seconds": 3600, + "limit": 1, + }, }, json={"next": 2, "items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]}, ) @@ -510,10 +528,9 @@ async def test_list_assets_pagination(monkeypatch, httpx_mock: pytest_httpx.HTTP url="https://api.armis.com/v3/assets/_search", method="POST", match_json={ - "limit": 1, "asset_type": "DEVICE", "fields": assets_test_data.ALL_DEVICE_FIELDS, - "filter": {"filter_criteria": "LAST_SEEN", "last_seen_seconds": 3600, "after": 2}, + "filter": {"filter_criteria": "LAST_SEEN", "last_seen_seconds": 3600, "limit": 1, "after": 2}, }, json={"next": None, "items": [{"asset_id": 2, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]}, )