|
3 | 3 | import logging |
4 | 4 | import os |
5 | 5 | import time |
| 6 | +import uuid |
6 | 7 | from datetime import datetime |
7 | 8 |
|
| 9 | +import pytest |
| 10 | + |
8 | 11 | from bedrock_agentcore.memory import MemoryClient |
9 | 12 |
|
10 | 13 | # Use INFO level logging for cleaner output |
@@ -408,5 +411,70 @@ def main(): |
408 | 411 | logger.error("Failed to delete test memory: %s", e) |
409 | 412 |
|
410 | 413 |
|
| 414 | +@pytest.mark.integration |
| 415 | +class TestMemoryClient: |
| 416 | + """Integration tests for streamDeliveryResources on MemoryClient.""" |
| 417 | + |
| 418 | + @classmethod |
| 419 | + def setup_class(cls): |
| 420 | + cls.kinesis_stream_arn = os.environ.get("MEMORY_KINESIS_ARN") |
| 421 | + cls.execution_role_arn = os.environ.get("MEMORY_ROLE_ARN") |
| 422 | + |
| 423 | + if not cls.kinesis_stream_arn or not cls.execution_role_arn: |
| 424 | + pytest.fail("MEMORY_KINESIS_ARN and MEMORY_ROLE_ARN must be set") |
| 425 | + |
| 426 | + cls.region = os.environ.get("BEDROCK_TEST_REGION", "us-west-2") |
| 427 | + cls.client = MemoryClient(region_name=cls.region) |
| 428 | + cls.test_prefix = f"test_stream_{int(time.time())}" |
| 429 | + cls.memory_ids = [] |
| 430 | + |
| 431 | + @classmethod |
| 432 | + def teardown_class(cls): |
| 433 | + for memory_id in cls.memory_ids: |
| 434 | + try: |
| 435 | + cls.client.delete_memory(memory_id) |
| 436 | + except Exception as e: |
| 437 | + print(f"Failed to delete memory {memory_id}: {e}") |
| 438 | + |
| 439 | + def _make_delivery_config(self, level="FULL_CONTENT"): |
| 440 | + return { |
| 441 | + "resources": [ |
| 442 | + { |
| 443 | + "kinesis": { |
| 444 | + "dataStreamArn": self.kinesis_stream_arn, |
| 445 | + "contentConfigurations": [{"type": "MEMORY_RECORDS", "level": level}], |
| 446 | + } |
| 447 | + } |
| 448 | + ] |
| 449 | + } |
| 450 | + |
| 451 | + def test_stream_delivery_create_and_update(self): |
| 452 | + """Create memory with FULL_CONTENT delivery, update to METADATA_ONLY via passthrough.""" |
| 453 | + delivery_config = self._make_delivery_config("FULL_CONTENT") |
| 454 | + |
| 455 | + memory = self.client.create_memory_and_wait( |
| 456 | + name=f"{self.test_prefix}_stream", |
| 457 | + strategies=[], |
| 458 | + memory_execution_role_arn=self.execution_role_arn, |
| 459 | + stream_delivery_resources=delivery_config, |
| 460 | + ) |
| 461 | + |
| 462 | + memory_id = memory["id"] |
| 463 | + self.__class__.memory_ids.append(memory_id) |
| 464 | + |
| 465 | + assert memory["streamDeliveryResources"] == delivery_config |
| 466 | + |
| 467 | + # Test update via MemoryClient.__getattr__ passthrough to boto3 client. |
| 468 | + # Uses camelCase params because the passthrough forwards directly to boto3 |
| 469 | + # without the snake_case translation that explicit SDK methods provide. |
| 470 | + updated_config = self._make_delivery_config("METADATA_ONLY") |
| 471 | + response = self.client.update_memory( |
| 472 | + memoryId=memory_id, |
| 473 | + clientToken=str(uuid.uuid4()), |
| 474 | + streamDeliveryResources=updated_config, |
| 475 | + ) |
| 476 | + assert response["memory"]["streamDeliveryResources"] == updated_config |
| 477 | + |
| 478 | + |
411 | 479 | if __name__ == "__main__": |
412 | 480 | main() |
0 commit comments