|
29 | 29 | format_response, |
30 | 30 | get_date_range, |
31 | 31 | handle_aws_error, |
| 32 | + paginate_aws_response, |
32 | 33 | parse_json, |
33 | 34 | validate_date_format, |
34 | 35 | ) |
@@ -396,3 +397,87 @@ def test_format_error_response_with_classification(self): |
396 | 397 | 'operation': 'getCostAndUsage', |
397 | 398 | 'service': 'Cost Explorer', |
398 | 399 | } |
| 400 | + |
| 401 | + |
| 402 | +class TestPaginateAwsResponse: |
| 403 | + """Tests for paginate_aws_response pagination metadata.""" |
| 404 | + |
| 405 | + @pytest.mark.asyncio |
| 406 | + async def test_last_page_empty_string_token(self): |
| 407 | + """Test that an empty-string token on the last page reports no more pages.""" |
| 408 | + # Setup |
| 409 | + ctx = AsyncMock() |
| 410 | + # Some AWS APIs (e.g. Savings Plans DescribeSavingsPlansOfferingRates) |
| 411 | + # return an empty-string token on the final page rather than omitting it. |
| 412 | + api_function = MagicMock(return_value={'Items': [1, 2, 3], 'NextToken': ''}) |
| 413 | + |
| 414 | + # Execute |
| 415 | + results, metadata = await paginate_aws_response( |
| 416 | + ctx, |
| 417 | + 'TestOperation', |
| 418 | + api_function, |
| 419 | + {}, |
| 420 | + 'Items', |
| 421 | + ) |
| 422 | + |
| 423 | + # Assert |
| 424 | + assert results == [1, 2, 3] |
| 425 | + assert metadata['has_more'] is False |
| 426 | + assert metadata['complete_dataset'] is True |
| 427 | + assert metadata['next_token'] is None |
| 428 | + assert metadata['pages_fetched'] == 1 |
| 429 | + assert metadata['total_results'] == 3 |
| 430 | + |
| 431 | + @pytest.mark.asyncio |
| 432 | + async def test_max_pages_hit_with_token(self): |
| 433 | + """Test that hitting max_pages with a real token still reports more pages.""" |
| 434 | + # Setup |
| 435 | + ctx = AsyncMock() |
| 436 | + api_function = MagicMock( |
| 437 | + side_effect=[ |
| 438 | + {'Items': [1, 2], 'NextToken': 'token-1'}, |
| 439 | + {'Items': [3, 4], 'NextToken': 'token-2'}, |
| 440 | + ] |
| 441 | + ) |
| 442 | + |
| 443 | + # Execute |
| 444 | + results, metadata = await paginate_aws_response( |
| 445 | + ctx, |
| 446 | + 'TestOperation', |
| 447 | + api_function, |
| 448 | + {}, |
| 449 | + 'Items', |
| 450 | + max_pages=2, |
| 451 | + ) |
| 452 | + |
| 453 | + # Assert |
| 454 | + assert results == [1, 2, 3, 4] |
| 455 | + assert metadata['has_more'] is True |
| 456 | + assert metadata['complete_dataset'] is False |
| 457 | + assert metadata['next_token'] == 'token-2' |
| 458 | + assert metadata['pages_fetched'] == 2 |
| 459 | + assert metadata['total_results'] == 4 |
| 460 | + |
| 461 | + @pytest.mark.asyncio |
| 462 | + async def test_missing_token_key(self): |
| 463 | + """Test that a missing token key reports no more pages (unchanged behavior).""" |
| 464 | + # Setup |
| 465 | + ctx = AsyncMock() |
| 466 | + api_function = MagicMock(return_value={'Items': [1, 2, 3]}) |
| 467 | + |
| 468 | + # Execute |
| 469 | + results, metadata = await paginate_aws_response( |
| 470 | + ctx, |
| 471 | + 'TestOperation', |
| 472 | + api_function, |
| 473 | + {}, |
| 474 | + 'Items', |
| 475 | + ) |
| 476 | + |
| 477 | + # Assert |
| 478 | + assert results == [1, 2, 3] |
| 479 | + assert metadata['has_more'] is False |
| 480 | + assert metadata['complete_dataset'] is True |
| 481 | + assert metadata['next_token'] is None |
| 482 | + assert metadata['pages_fetched'] == 1 |
| 483 | + assert metadata['total_results'] == 3 |
0 commit comments