Skip to content

Commit f6b9bc0

Browse files
read-write-content[bot]read-write-content[bot]
authored
Pass limit parameter in the correct location during assets pagination (#58)
1 parent 49c5139 commit f6b9bc0

4 files changed

Lines changed: 51 additions & 25 deletions

File tree

‎armis_sdk/clients/assets_client.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -501,7 +501,7 @@ async def _list_assets(
501501
"fields": fields,
502502
"filter": filter_,
503503
}
504-
async for item in self._armis_client.list("/v3/assets/_search", body=body, after_location="filter"):
504+
async for item in self._armis_client.list("/v3/assets/_search", body=body, pagination_location="filter"):
505505
yield asset_class.from_search_result(item)
506506

507507
@classmethod

‎armis_sdk/core/armis_client.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -89,13 +89,18 @@ def client(self, retries: int | None = None, backoff: float | None = None):
8989
trust_env=True,
9090
)
9191

92-
async def list(self, url: str, body: dict | None = None, after_location: str | None = None) -> AsyncIterator[dict]:
92+
async def list(
93+
self,
94+
url: str,
95+
body: dict | None = None,
96+
pagination_location: str | None = None,
97+
) -> AsyncIterator[dict]:
9398
"""List all items from a paginated endpoint.
9499
95100
Args:
96101
url (str): The relative endpoint URL.
97102
body (dict): Payload to send as POST request.
98-
after_location (str): The nested object location to use for pagination.
103+
pagination_location (str): The nested object location to use for pagination.
99104
100105
Returns:
101106
An (async) iterator of `dict`s.
@@ -123,7 +128,11 @@ async def main():
123128
"""
124129
page_size = int(os.getenv(ARMIS_PAGE_SIZE, str(DEFAULT_PAGE_LENGTH)))
125130
async with self.client() as client:
126-
params = {"limit": page_size, **(body or {})}
131+
params = {**(body or {})}
132+
if pagination_location:
133+
params[pagination_location]["limit"] = page_size
134+
else:
135+
params["limit"] = page_size
127136
while True:
128137
if body:
129138
response = await client.post(url, json=params)
@@ -134,8 +143,8 @@ async def main():
134143
for item in items:
135144
yield item
136145
if next_ := data.get("next"):
137-
if after_location:
138-
params[after_location]["after"] = next_
146+
if pagination_location:
147+
params[pagination_location]["after"] = next_
139148
else:
140149
params["after"] = next_
141150
else:

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "armis_sdk"
3-
version = "1.2.1"
3+
version = "1.2.2"
44
description = "The Armis SDK is a package that encapsulates common use-cases for interacting with the Armis platform."
55
authors = [
66
{ name = "Shai Lachmanovich", email = "shai@armis.com" },

‎tests/armis_sdk/clients/assets_client_test.py‎

Lines changed: 35 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,12 @@ async def test_list_by_last_seen_datetime(httpx_mock: pytest_httpx.HTTPXMock):
2525
url="https://api.armis.com/v3/assets/_search",
2626
method="POST",
2727
match_json={
28-
"limit": 100,
2928
"asset_type": "DEVICE",
3029
"fields": assets_test_data.ALL_DEVICE_FIELDS,
3130
"filter": {
3231
"filter_criteria": "LAST_SEEN",
3332
"last_seen_ge": "2025-12-03T00:00:00",
33+
"limit": 100,
3434
},
3535
},
3636
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(
5050
url="https://api.armis.com/v3/assets/_search",
5151
method="POST",
5252
match_json={
53-
"limit": 100,
5453
"asset_type": "DEVICE",
5554
"fields": ["brand", "custom.MyField1", "custom.MyField2", "purdue_level"],
5655
"filter": {
5756
"filter_criteria": "LAST_SEEN",
5857
"last_seen_ge": "2025-12-03T00:00:00",
58+
"limit": 100,
5959
},
6060
},
6161
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):
7474
url="https://api.armis.com/v3/assets/_search",
7575
method="POST",
7676
match_json={
77-
"limit": 100,
7877
"asset_type": "DEVICE",
7978
"fields": assets_test_data.ALL_DEVICE_FIELDS,
80-
"filter": {"filter_criteria": "LAST_SEEN", "last_seen_seconds": 3600},
79+
"filter": {
80+
"filter_criteria": "LAST_SEEN",
81+
"last_seen_seconds": 3600,
82+
"limit": 100,
83+
},
8184
},
8285
json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]},
8386
)
@@ -96,10 +99,13 @@ async def test_list_by_last_seen_timedelta_explicit_fields(
9699
url="https://api.armis.com/v3/assets/_search",
97100
method="POST",
98101
match_json={
99-
"limit": 100,
100102
"asset_type": "DEVICE",
101103
"fields": ["brand", "custom.MyField1", "custom.MyField2", "purdue_level"],
102-
"filter": {"filter_criteria": "LAST_SEEN", "last_seen_seconds": 3600},
104+
"filter": {
105+
"filter_criteria": "LAST_SEEN",
106+
"last_seen_seconds": 3600,
107+
"limit": 100,
108+
},
103109
},
104110
json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_PARTIAL_RAW_DATA}]},
105111
)
@@ -130,13 +136,13 @@ async def test_list_by_asset_id(httpx_mock: pytest_httpx.HTTPXMock):
130136
url="https://api.armis.com/v3/assets/_search",
131137
method="POST",
132138
match_json={
133-
"limit": 100,
134139
"asset_type": "DEVICE",
135140
"fields": assets_test_data.ALL_DEVICE_FIELDS,
136141
"filter": {
137142
"filter_criteria": "ASSET_ID",
138143
"asset_id_source": "IPV4_ADDRESS",
139144
"asset_ids": ["1.1.1.1"],
145+
"limit": 100,
140146
},
141147
},
142148
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
161167
url="https://api.armis.com/v3/assets/_search",
162168
method="POST",
163169
match_json={
164-
"limit": 100,
165170
"asset_type": "DEVICE",
166171
"fields": ["brand", "custom.MyField1", "custom.MyField2", "purdue_level"],
167172
"filter": {
168173
"filter_criteria": "ASSET_ID",
169174
"asset_id_source": "IPV4_ADDRESS",
170175
"asset_ids": ["1.1.1.1"],
176+
"limit": 100,
171177
},
172178
},
173179
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):
205211
url="https://api.armis.com/v3/assets/_search",
206212
method="POST",
207213
match_json={
208-
"limit": 100,
209214
"asset_type": "DEVICE",
210215
"fields": assets_test_data.ALL_DEVICE_FIELDS,
211-
"filter": {"filter_criteria": "BOUNDARY_ID", "boundary_ids": [1, 2, 3]},
216+
"filter": {
217+
"filter_criteria": "BOUNDARY_ID",
218+
"boundary_ids": [1, 2, 3],
219+
"limit": 100,
220+
},
212221
},
213222
json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]},
214223
)
@@ -253,10 +262,13 @@ async def test_list_by_multiple(httpx_mock: pytest_httpx.HTTPXMock, kwargs, expe
253262
url="https://api.armis.com/v3/assets/_search",
254263
method="POST",
255264
match_json={
256-
"limit": 100,
257265
"asset_type": "DEVICE",
258266
"fields": assets_test_data.ALL_DEVICE_FIELDS,
259-
"filter": {"filter_criteria": "MULTIPLE", "filters": expected_filters},
267+
"filter": {
268+
"filter_criteria": "MULTIPLE",
269+
"filters": expected_filters,
270+
"limit": 100,
271+
},
260272
},
261273
json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]},
262274
)
@@ -313,10 +325,13 @@ async def test_list_by_site_id(httpx_mock: pytest_httpx.HTTPXMock):
313325
url="https://api.armis.com/v3/assets/_search",
314326
method="POST",
315327
match_json={
316-
"limit": 100,
317328
"asset_type": "DEVICE",
318329
"fields": assets_test_data.ALL_DEVICE_FIELDS,
319-
"filter": {"filter_criteria": "SITE_ID", "site_ids": [1, 2, 3]},
330+
"filter": {
331+
"filter_criteria": "SITE_ID",
332+
"site_ids": [1, 2, 3],
333+
"limit": 100,
334+
},
320335
},
321336
json={"items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]},
322337
)
@@ -499,21 +514,23 @@ async def test_list_assets_pagination(monkeypatch, httpx_mock: pytest_httpx.HTTP
499514
url="https://api.armis.com/v3/assets/_search",
500515
method="POST",
501516
match_json={
502-
"limit": 1,
503517
"asset_type": "DEVICE",
504518
"fields": assets_test_data.ALL_DEVICE_FIELDS,
505-
"filter": {"filter_criteria": "LAST_SEEN", "last_seen_seconds": 3600},
519+
"filter": {
520+
"filter_criteria": "LAST_SEEN",
521+
"last_seen_seconds": 3600,
522+
"limit": 1,
523+
},
506524
},
507525
json={"next": 2, "items": [{"asset_id": 1, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]},
508526
)
509527
httpx_mock.add_response(
510528
url="https://api.armis.com/v3/assets/_search",
511529
method="POST",
512530
match_json={
513-
"limit": 1,
514531
"asset_type": "DEVICE",
515532
"fields": assets_test_data.ALL_DEVICE_FIELDS,
516-
"filter": {"filter_criteria": "LAST_SEEN", "last_seen_seconds": 3600, "after": 2},
533+
"filter": {"filter_criteria": "LAST_SEEN", "last_seen_seconds": 3600, "limit": 1, "after": 2},
517534
},
518535
json={"next": None, "items": [{"asset_id": 2, "fields": assets_test_data.MOCK_DEVICE_FULL_RAW_DATA}]},
519536
)

0 commit comments

Comments
 (0)