1
1
# SPDX-License-Identifier: Apache-2.0
2
2
# Copyright 2022 Atlan Pte. Ltd.
3
+ import math
3
4
from dataclasses import dataclass , field
4
5
from datetime import datetime
5
6
from time import sleep , time
@@ -335,8 +336,6 @@ def _assert_search_results(results, expected_sorts, size, TOTAL_ASSETS, bulk=Fal
335
336
336
337
@patch .object (LOGGER , "debug" )
337
338
def test_search_pagination (mock_logger , client : AtlanClient ):
338
- size = 2
339
-
340
339
# Avoid testing on integration tests objects
341
340
exclude_sdk_terms = [
342
341
Asset .NAME .wildcard ("psdk_*" ),
@@ -352,13 +351,22 @@ def test_search_pagination(mock_logger, client: AtlanClient):
352
351
dsl = DSL (
353
352
query = query ,
354
353
post_filter = Term .with_type_name (value = "AtlasGlossaryTerm" ),
355
- size = size ,
354
+ size = 0 , # to get the total count
356
355
)
356
+
357
357
request = IndexSearchRequest (dsl = dsl )
358
358
results = client .asset .search (criteria = request )
359
359
# Assigning this here to ensure the total assets
360
360
# remain constant across different test cases
361
361
TOTAL_ASSETS = results .count
362
+
363
+ # set page_size to divide into ~5 API calls
364
+ size = max (1 , math .ceil (TOTAL_ASSETS / 5 ))
365
+ request .dsl .size = size
366
+
367
+ # Now, we can test different test scenarios for search() with the dynamic page size
368
+ results = client .asset .search (criteria = request )
369
+
362
370
expected_sorts = [Asset .GUID .order (SortOrder .ASCENDING )]
363
371
_assert_search_results (results , expected_sorts , size , TOTAL_ASSETS )
364
372
@@ -385,7 +393,7 @@ def test_search_pagination(mock_logger, client: AtlanClient):
385
393
FluentSearch (where_nots = exclude_sdk_terms )
386
394
.where (CompoundQuery .active_assets ())
387
395
.where (CompoundQuery .asset_type (AtlasGlossaryTerm ))
388
- .page_size (2 )
396
+ .page_size (size )
389
397
).to_request ()
390
398
results = client .asset .search (criteria = request )
391
399
expected_sorts = [Asset .GUID .order (SortOrder .ASCENDING )]
@@ -396,7 +404,7 @@ def test_search_pagination(mock_logger, client: AtlanClient):
396
404
FluentSearch (where_nots = exclude_sdk_terms )
397
405
.where (CompoundQuery .active_assets ())
398
406
.where (CompoundQuery .asset_type (AtlasGlossaryTerm ))
399
- .page_size (2 )
407
+ .page_size (size )
400
408
).to_request ()
401
409
results = client .asset .search (criteria = request , bulk = True )
402
410
expected_sorts = [
@@ -413,7 +421,7 @@ def test_search_pagination(mock_logger, client: AtlanClient):
413
421
FluentSearch (where_nots = exclude_sdk_terms )
414
422
.where (CompoundQuery .active_assets ())
415
423
.where (CompoundQuery .asset_type (AtlasGlossaryTerm ))
416
- .page_size (2 )
424
+ .page_size (size )
417
425
).execute (client , bulk = True )
418
426
expected_sorts = [
419
427
Asset .CREATE_TIME .order (SortOrder .ASCENDING ),
@@ -431,7 +439,7 @@ def test_search_pagination(mock_logger, client: AtlanClient):
431
439
FluentSearch (where_nots = exclude_sdk_terms )
432
440
.where (CompoundQuery .active_assets ())
433
441
.where (CompoundQuery .asset_type (AtlasGlossaryTerm ))
434
- .page_size (2 )
442
+ .page_size (size )
435
443
).to_request ()
436
444
results = client .asset .search (criteria = request )
437
445
expected_sorts = [
0 commit comments