diff --git a/src/bedrock_agentcore/memory/controlplane.py b/src/bedrock_agentcore/memory/controlplane.py index 25e75d0..c6f7026 100644 --- a/src/bedrock_agentcore/memory/controlplane.py +++ b/src/bedrock_agentcore/memory/controlplane.py @@ -1,44 +1,68 @@ -"""AgentCore Memory SDK - Control Plane Client. +"""AgentCore Memory SDK - Control Plane Client (DEPRECATED). -This module provides a simplified interface for Bedrock AgentCore Memory control plane operations. -It handles memory resource management, strategy operations, and status monitoring. +This module is deprecated. Use MemoryClient from bedrock_agentcore.memory.client instead, +which provides all control plane operations plus data plane features. """ import logging import os import time import uuid +import warnings from typing import Any, Dict, List, Optional import boto3 from botocore.exceptions import ClientError -from .constants import ( - MemoryStatus, -) +from .constants import MemoryStatus logger = logging.getLogger(__name__) +def _deprecation(method_name: str, suggestion: str) -> None: + """Emit a deprecation warning for a MemoryControlPlaneClient method.""" + warnings.warn( + f"MemoryControlPlaneClient.{method_name}() is deprecated. Use MemoryClient.{suggestion} instead.", + DeprecationWarning, + stacklevel=3, + ) + + class MemoryControlPlaneClient: - """Client for Bedrock AgentCore Memory control plane operations.""" + """Client for Bedrock AgentCore Memory control plane operations. + + .. deprecated:: + Use :class:`MemoryClient` instead, which provides all control plane operations + plus data plane features. + """ - def __init__(self, region_name: str = "us-west-2", environment: str = "prod"): + def __init__(self, region_name: str = "us-west-2", environment: str = "prod", client=None): """Initialize the Memory Control Plane client. Args: region_name: AWS region name - environment: Environment name (prod, gamma, etc.) + environment: Environment name (unused, kept for backward compatibility) + client: Optional pre-built boto3 client (for testing/backward compatibility) """ + warnings.warn( + "MemoryControlPlaneClient is deprecated and will be removed in v1.4.0. " + "Use MemoryClient instead, which provides all control plane operations " + "plus data plane features.", + DeprecationWarning, + stacklevel=2, + ) self.region_name = region_name self.environment = environment - self.endpoint = os.getenv( - "BEDROCK_AGENTCORE_CONTROL_ENDPOINT", f"https://bedrock-agentcore-control.{region_name}.amazonaws.com" - ) - - service_name = os.getenv("BEDROCK_AGENTCORE_CONTROL_SERVICE", "bedrock-agentcore-control") - self.client = boto3.client(service_name, region_name=self.region_name, endpoint_url=self.endpoint) + if client is not None: + self.client = client + else: + self.endpoint = os.getenv( + "BEDROCK_AGENTCORE_CONTROL_ENDPOINT", + f"https://bedrock-agentcore-control.{region_name}.amazonaws.com", + ) + service_name = os.getenv("BEDROCK_AGENTCORE_CONTROL_SERVICE", "bedrock-agentcore-control") + self.client = boto3.client(service_name, region_name=self.region_name, endpoint_url=self.endpoint) logger.info("Initialized MemoryControlPlaneClient for %s in %s", environment, region_name) @@ -57,6 +81,9 @@ def create_memory( ) -> Dict[str, Any]: """Create a memory resource with optional strategies. + .. deprecated:: + Use :meth:`MemoryClient.create_memory` or :meth:`MemoryClient.create_memory_and_wait`. + Args: name: Name for the memory resource event_expiry_days: How long to retain events (default: 90 days) @@ -70,6 +97,8 @@ def create_memory( Returns: Created memory object """ + _deprecation("create_memory", "create_memory() or create_memory_and_wait()") + params = { "name": name, "eventExpiryDuration": event_expiry_days, @@ -104,6 +133,11 @@ def create_memory( def get_memory(self, memory_id: str, include_strategies: bool = True) -> Dict[str, Any]: """Get a memory resource by ID. + .. deprecated:: + Use ``MemoryClient.get_memory(memoryId=memory_id)`` via the client's + ``__getattr__`` forwarding, or :meth:`MemoryClient.get_memory_strategies` + for strategy details. + Args: memory_id: Memory resource ID include_strategies: Whether to include strategy details in response @@ -111,6 +145,8 @@ def get_memory(self, memory_id: str, include_strategies: bool = True) -> Dict[st Returns: Memory resource details """ + _deprecation("get_memory", "get_memory(memoryId=...) or get_memory_strategies()") + try: response = self.client.get_memory(memoryId=memory_id) memory = response["memory"] @@ -132,12 +168,17 @@ def get_memory(self, memory_id: str, include_strategies: bool = True) -> Dict[st def list_memories(self, max_results: int = 100) -> List[Dict[str, Any]]: """List all memories for the account with pagination support. + .. deprecated:: + Use :meth:`MemoryClient.list_memories`. + Args: max_results: Maximum number of memories to return Returns: List of memory summaries """ + _deprecation("list_memories", "list_memories()") + try: memories = [] next_token = None @@ -180,6 +221,10 @@ def update_memory( ) -> Dict[str, Any]: """Update a memory resource properties and/or strategies. + .. deprecated:: + Use :meth:`MemoryClient.update_memory_strategies` for strategy changes, + or call ``MemoryClient.update_memory(...)`` directly for property updates. + Args: memory_id: Memory resource ID description: Optional new description @@ -195,6 +240,11 @@ def update_memory( Returns: Updated memory object """ + _deprecation( + "update_memory", + "update_memory_strategies() or update_memory_strategies_and_wait()", + ) + params: Dict = { "memoryId": memory_id, "clientToken": str(uuid.uuid4()), @@ -245,12 +295,15 @@ def delete_memory( self, memory_id: str, wait_for_deletion: bool = False, - wait_for_strategies: bool = False, # Changed default to False + wait_for_strategies: bool = False, max_wait: int = 300, poll_interval: int = 10, ) -> Dict[str, Any]: """Delete a memory resource. + .. deprecated:: + Use :meth:`MemoryClient.delete_memory` or :meth:`MemoryClient.delete_memory_and_wait`. + Args: memory_id: Memory resource ID to delete wait_for_deletion: Whether to wait for complete deletion @@ -261,6 +314,8 @@ def delete_memory( Returns: Deletion response """ + _deprecation("delete_memory", "delete_memory() or delete_memory_and_wait()") + try: # If requested, wait for all strategies to become ACTIVE before deletion if wait_for_strategies: @@ -277,7 +332,8 @@ def delete_memory( if transitional_strategies: logger.info( - "Waiting for %d strategies to become ACTIVE before deletion", len(transitional_strategies) + "Waiting for %d strategies to become ACTIVE before deletion", + len(transitional_strategies), ) self._wait_for_status( memory_id=memory_id, @@ -327,6 +383,10 @@ def add_strategy( ) -> Dict[str, Any]: """Add a strategy to a memory resource. + .. deprecated:: + Use :meth:`MemoryClient.update_memory_strategies` with ``add_strategies``, + or one of the typed helpers like :meth:`MemoryClient.add_semantic_strategy`. + Args: memory_id: Memory resource ID strategy: Strategy configuration dictionary @@ -337,6 +397,11 @@ def add_strategy( Returns: Updated memory object with strategyId field """ + _deprecation( + "add_strategy", + "update_memory_strategies(add_strategies=[...]) or add_semantic_strategy() / add_summary_strategy()", + ) + # Get the strategy type and name for identification strategy_type = list(strategy.keys())[0] # e.g., 'semanticMemoryStrategy' strategy_name = strategy[strategy_type].get("name") @@ -376,6 +441,9 @@ def add_strategy( def get_strategy(self, memory_id: str, strategy_id: str) -> Dict[str, Any]: """Get a specific strategy from a memory resource. + .. deprecated:: + Use :meth:`MemoryClient.get_memory_strategies` and filter by strategy ID. + Args: memory_id: Memory resource ID strategy_id: Strategy ID @@ -383,6 +451,8 @@ def get_strategy(self, memory_id: str, strategy_id: str) -> Dict[str, Any]: Returns: Strategy details """ + _deprecation("get_strategy", "get_memory_strategies()") + try: memory = self.get_memory(memory_id) strategies = memory.get("strategies", []) @@ -410,6 +480,10 @@ def update_strategy( ) -> Dict[str, Any]: """Update a strategy in a memory resource. + .. deprecated:: + Use :meth:`MemoryClient.update_memory_strategies` with ``modify_strategies``, + or :meth:`MemoryClient.modify_strategy`. + Args: memory_id: Memory resource ID strategy_id: Strategy ID to update @@ -423,6 +497,11 @@ def update_strategy( Returns: Updated memory object """ + _deprecation( + "update_strategy", + "update_memory_strategies(modify_strategies=[...]) or modify_strategy()", + ) + # Note: API expects memoryStrategyId for input but returns strategyId in response modify_config: Dict = {"memoryStrategyId": strategy_id} @@ -458,6 +537,10 @@ def remove_strategy( ) -> Dict[str, Any]: """Remove a strategy from a memory resource. + .. deprecated:: + Use :meth:`MemoryClient.update_memory_strategies` with ``delete_strategy_ids``, + or :meth:`MemoryClient.delete_strategy`. + Args: memory_id: Memory resource ID strategy_id: Strategy ID to remove @@ -468,6 +551,11 @@ def remove_strategy( Returns: Updated memory object """ + _deprecation( + "remove_strategy", + "update_memory_strategies(delete_strategy_ids=[...]) or delete_strategy()", + ) + # For remove_strategy, we only need to wait for memory to be active # since the strategy will be gone return self.update_memory( @@ -484,7 +572,10 @@ def _wait_for_memory_active(self, memory_id: str, max_wait: int, poll_interval: """Wait for memory to return to ACTIVE state.""" logger.info("Waiting for memory %s to become ACTIVE...", memory_id) return self._wait_for_status( - memory_id=memory_id, target_status=MemoryStatus.ACTIVE.value, max_wait=max_wait, poll_interval=poll_interval + memory_id=memory_id, + target_status=MemoryStatus.ACTIVE.value, + max_wait=max_wait, + poll_interval=poll_interval, ) def _wait_for_strategy_active( @@ -535,7 +626,12 @@ def _wait_for_strategy_active( ) def _wait_for_status( - self, memory_id: str, target_status: str, max_wait: int, poll_interval: int, check_strategies: bool = True + self, + memory_id: str, + target_status: str, + max_wait: int, + poll_interval: int, + check_strategies: bool = True, ) -> Dict[str, Any]: """Generic method to wait for a memory to reach a specific status. @@ -606,7 +702,10 @@ def _wait_for_status( elapsed = time.time() - start_time logger.info( - "Memory %s and all strategies are now %s (took %.1f seconds)", memory_id, target_status, elapsed + "Memory %s and all strategies are now %s (took %.1f seconds)", + memory_id, + target_status, + elapsed, ) return memory elif status == MemoryStatus.FAILED.value: diff --git a/tests/bedrock_agentcore/memory/test_controlplane.py b/tests/bedrock_agentcore/memory/test_controlplane.py index 38b691c..bcc6284 100644 --- a/tests/bedrock_agentcore/memory/test_controlplane.py +++ b/tests/bedrock_agentcore/memory/test_controlplane.py @@ -1,792 +1,562 @@ """Unit tests for Memory Control Plane Client - no external connections.""" import uuid +import warnings from unittest.mock import MagicMock, patch +import pytest from botocore.exceptions import ClientError from bedrock_agentcore.memory.constants import MemoryStatus from bedrock_agentcore.memory.controlplane import MemoryControlPlaneClient +# Suppress the deprecation warnings for all tests except the ones that explicitly test them +pytestmark = pytest.mark.filterwarnings("ignore::DeprecationWarning") + + +def _create_client(): + """Helper to create a MemoryControlPlaneClient with a mocked boto3 client passed via constructor.""" + mock_client = MagicMock() + client = MemoryControlPlaneClient(client=mock_client) + return client, mock_client + + +def test_deprecation_warning(): + """Test that MemoryControlPlaneClient emits a deprecation warning on init.""" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + MemoryControlPlaneClient(client=MagicMock()) + deprecation_warnings = [x for x in w if issubclass(x.category, DeprecationWarning)] + assert len(deprecation_warnings) >= 1 + assert "MemoryClient" in str(deprecation_warnings[0].message) + assert "v1.4.0" in str(deprecation_warnings[0].message) + + +def test_method_deprecation_warnings(): + """Test that each public method emits its own deprecation warning.""" + client, mock_client = _create_client() + + mock_client.create_memory.return_value = {"memory": {"id": "mem-123", "name": "Test", "status": "CREATING"}} + mock_client.get_memory.return_value = {"memory": {"id": "mem-123", "status": "ACTIVE", "strategies": []}} + mock_client.list_memories.return_value = {"memories": [], "nextToken": None} + mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} + mock_client.delete_memory.return_value = {"status": "DELETING"} + + methods_and_suggestions = [ + (lambda: client.create_memory(name="Test"), "create_memory"), + (lambda: client.get_memory("mem-123"), "get_memory"), + (lambda: client.list_memories(), "list_memories"), + (lambda: client.update_memory(memory_id="mem-123"), "update_memory"), + (lambda: client.delete_memory("mem-123"), "delete_memory"), + (lambda: client.get_strategy("mem-123", "strat-1"), "get_strategy"), + ] + + for method_call, method_name in methods_and_suggestions: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + try: + method_call() + except (ValueError, ClientError): + pass # Some methods may raise due to mock setup + deprecation_warnings = [x for x in w if issubclass(x.category, DeprecationWarning)] + method_warnings = [ + x for x in deprecation_warnings if f"MemoryControlPlaneClient.{method_name}()" in str(x.message) + ] + assert len(method_warnings) >= 1, f"No deprecation warning for {method_name}" + assert "MemoryClient" in str(method_warnings[0].message) + def test_create_memory(): """Test create_memory functionality.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # Mock successful response - mock_client.create_memory.return_value = { - "memory": {"id": "mem-123", "name": "Test Memory", "status": "CREATING"} - } + # Mock successful response + mock_client.create_memory.return_value = {"memory": {"id": "mem-123", "name": "Test Memory", "status": "CREATING"}} - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - # Test basic memory creation - result = client.create_memory(name="Test Memory", description="Test description") + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + result = client.create_memory(name="Test Memory", description="Test description") - assert result["id"] == "mem-123" - assert result["name"] == "Test Memory" - assert mock_client.create_memory.called + assert result["id"] == "mem-123" + assert result["name"] == "Test Memory" + assert mock_client.create_memory.called - # Verify correct parameters were passed - args, kwargs = mock_client.create_memory.call_args - assert kwargs["name"] == "Test Memory" - assert kwargs["description"] == "Test description" - assert kwargs["clientToken"] == "12345678-1234-5678-1234-567812345678" + # Verify core parameters were passed + args, kwargs = mock_client.create_memory.call_args + assert kwargs["name"] == "Test Memory" + assert kwargs["description"] == "Test description" + assert kwargs["clientToken"] == "12345678-1234-5678-1234-567812345678" def test_get_memory(): """Test get_memory functionality.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # Mock response with strategies - mock_client.get_memory.return_value = { - "memory": { - "id": "mem-123", - "name": "Test Memory", - "status": "ACTIVE", - "strategies": [ - {"strategyId": "strat-1", "type": "SEMANTIC"}, - {"strategyId": "strat-2", "type": "SUMMARY"}, - ], - } + client, mock_client = _create_client() + + # Mock response with strategies + mock_client.get_memory.return_value = { + "memory": { + "id": "mem-123", + "name": "Test Memory", + "status": "ACTIVE", + "strategies": [ + {"strategyId": "strat-1", "type": "SEMANTIC"}, + {"strategyId": "strat-2", "type": "SUMMARY"}, + ], } + } - # Test get memory with strategies - result = client.get_memory("mem-123") + result = client.get_memory("mem-123") - assert result["id"] == "mem-123" - assert result["strategyCount"] == 2 - assert "strategies" in result + assert result["id"] == "mem-123" + assert result["strategyCount"] == 2 + assert "strategies" in result - # Verify API call - mock_client.get_memory.assert_called_with(memoryId="mem-123") + # Verify API call + mock_client.get_memory.assert_called_with(memoryId="mem-123") def test_list_memories(): """Test list_memories functionality.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + mock_memories = [ + {"id": "mem-1", "name": "Memory 1", "status": "ACTIVE"}, + {"id": "mem-2", "name": "Memory 2", "status": "ACTIVE"}, + ] + mock_client.list_memories.return_value = {"memories": mock_memories, "nextToken": None} - # Mock response - mock_memories = [ - {"id": "mem-1", "name": "Memory 1", "status": "ACTIVE"}, - {"id": "mem-2", "name": "Memory 2", "status": "ACTIVE"}, - ] - mock_client.list_memories.return_value = {"memories": mock_memories, "nextToken": None} + result = client.list_memories(max_results=50) - # Test list memories - result = client.list_memories(max_results=50) + assert len(result) == 2 + assert result[0]["id"] == "mem-1" + assert result[0]["strategyCount"] == 0 + assert result[1]["strategyCount"] == 0 - assert len(result) == 2 - assert result[0]["id"] == "mem-1" - assert result[0]["strategyCount"] == 0 # List doesn't include strategies - - # Verify API call - args, kwargs = mock_client.list_memories.call_args - assert kwargs["maxResults"] == 50 + # Verify API call + args, kwargs = mock_client.list_memories.call_args + assert kwargs["maxResults"] == 50 def test_update_memory(): """Test update_memory functionality.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + mock_client.update_memory.return_value = { + "memory": {"id": "mem-123", "name": "Updated Memory", "status": "CREATING"} + } - # Mock response - mock_client.update_memory.return_value = { - "memory": {"id": "mem-123", "name": "Updated Memory", "status": "CREATING"} - } + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + result = client.update_memory(memory_id="mem-123", description="Updated description", event_expiry_days=120) - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - # Test memory update - result = client.update_memory(memory_id="mem-123", description="Updated description", event_expiry_days=120) - - assert result["id"] == "mem-123" - assert mock_client.update_memory.called + assert result["id"] == "mem-123" + assert mock_client.update_memory.called - # Verify correct parameters - args, kwargs = mock_client.update_memory.call_args - assert kwargs["memoryId"] == "mem-123" - assert kwargs["description"] == "Updated description" - assert kwargs["eventExpiryDuration"] == 120 - assert kwargs["clientToken"] == "12345678-1234-5678-1234-567812345678" + args, kwargs = mock_client.update_memory.call_args + assert kwargs["memoryId"] == "mem-123" + assert kwargs["description"] == "Updated description" + assert kwargs["eventExpiryDuration"] == 120 + assert kwargs["clientToken"] == "12345678-1234-5678-1234-567812345678" def test_delete_memory(): """Test delete_memory functionality.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + mock_client.delete_memory.return_value = {"status": "DELETING"} - # Mock response - mock_client.delete_memory.return_value = {"status": "DELETING"} + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + result = client.delete_memory("mem-123") - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - # Test memory deletion - result = client.delete_memory("mem-123") + assert result["status"] == "DELETING" + assert mock_client.delete_memory.called - assert result["status"] == "DELETING" - assert mock_client.delete_memory.called - - # Verify correct parameters - args, kwargs = mock_client.delete_memory.call_args - assert kwargs["memoryId"] == "mem-123" - assert kwargs["clientToken"] == "12345678-1234-5678-1234-567812345678" + args, kwargs = mock_client.delete_memory.call_args + assert kwargs["memoryId"] == "mem-123" + assert kwargs["clientToken"] == "12345678-1234-5678-1234-567812345678" def test_delete_memory_wait_for_strategies(): """Test delete_memory with wait_for_strategies=True.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # Mock get_memory response with strategies in transitional state - mock_client.get_memory.return_value = { + # First call from get_memory (initial check): one strategy still CREATING + # Second call from _wait_for_status polling: all active + mock_client.get_memory.side_effect = [ + { "memory": { "id": "mem-123", + "status": "ACTIVE", "strategies": [ - {"strategyId": "strat-1", "status": "CREATING"}, # Transitional state - {"strategyId": "strat-2", "status": "ACTIVE"}, # Already active + {"strategyId": "strat-1", "status": "CREATING"}, + {"strategyId": "strat-2", "status": "ACTIVE"}, ], } - } - - # Mock delete_memory response - mock_client.delete_memory.return_value = {"status": "DELETING"} - - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - with patch("time.time", return_value=0): - with patch("time.sleep"): - # Mock the _wait_for_status method to avoid actual waiting - with patch.object(client, "_wait_for_status") as mock_wait: - mock_wait.return_value = {"id": "mem-123", "status": "ACTIVE"} - - # Test memory deletion with wait_for_strategies=True - result = client.delete_memory("mem-123", wait_for_strategies=True) - - assert result["status"] == "DELETING" + }, + { + "memory": { + "id": "mem-123", + "status": "ACTIVE", + "strategies": [ + {"strategyId": "strat-1", "status": "ACTIVE"}, + {"strategyId": "strat-2", "status": "ACTIVE"}, + ], + } + }, + ] - # Verify get_memory was called to check strategy status - assert mock_client.get_memory.called + mock_client.delete_memory.return_value = {"status": "DELETING"} - # Verify _wait_for_status was called due to transitional strategy - mock_wait.assert_called_once_with( - memory_id="mem-123", - target_status=MemoryStatus.ACTIVE.value, - max_wait=300, - poll_interval=10, - check_strategies=True, - ) + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + with patch("time.time", return_value=0): + with patch("time.sleep"): + result = client.delete_memory("mem-123", wait_for_strategies=True) - # Verify delete_memory was called - assert mock_client.delete_memory.called - args, kwargs = mock_client.delete_memory.call_args - assert kwargs["memoryId"] == "mem-123" + assert result["status"] == "DELETING" + assert mock_client.get_memory.call_count == 2 + assert mock_client.delete_memory.called def test_delete_memory_wait_for_deletion(): """Test delete_memory with wait_for_deletion=True.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + client, mock_client = _create_client() - # Mock delete_memory response - mock_client.delete_memory.return_value = {"status": "DELETING"} + mock_client.delete_memory.return_value = {"status": "DELETING"} - # Mock get_memory to first return the memory, then raise ResourceNotFoundException - error_response = {"Error": {"Code": "ResourceNotFoundException", "Message": "Memory not found"}} - mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") + error_response = {"Error": {"Code": "ResourceNotFoundException", "Message": "Memory not found"}} + mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - with patch("time.time", return_value=0): - with patch("time.sleep"): - # Test memory deletion with wait_for_deletion=True - result = client.delete_memory("mem-123", wait_for_deletion=True, max_wait=120, poll_interval=5) - - assert result["status"] == "DELETING" - - # Verify delete_memory was called - assert mock_client.delete_memory.called - delete_args, delete_kwargs = mock_client.delete_memory.call_args - assert delete_kwargs["memoryId"] == "mem-123" - assert delete_kwargs["clientToken"] == "12345678-1234-5678-1234-567812345678" + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + with patch("time.time", return_value=0): + with patch("time.sleep"): + result = client.delete_memory("mem-123", wait_for_deletion=True, max_wait=120, poll_interval=5) - # Verify get_memory was called (to check if memory is gone) - assert mock_client.get_memory.called - get_args, get_kwargs = mock_client.get_memory.call_args - assert get_kwargs["memoryId"] == "mem-123" + assert result["status"] == "DELETING" + assert mock_client.delete_memory.called + assert mock_client.get_memory.called def test_add_strategy(): """Test add_strategy functionality.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + client, mock_client = _create_client() - # Mock update_memory response (add_strategy uses update_memory internally) - mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} + mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} - # Test strategy addition - strategy = {"semanticMemoryStrategy": {"name": "Test Strategy"}} + strategy = {"semanticMemoryStrategy": {"name": "Test Strategy"}} - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - result = client.add_strategy("mem-123", strategy) + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + result = client.add_strategy("mem-123", strategy) - assert result["id"] == "mem-123" - assert mock_client.update_memory.called + assert result["id"] == "mem-123" + assert mock_client.update_memory.called - # Verify strategy was passed correctly - args, kwargs = mock_client.update_memory.call_args - assert "memoryStrategies" in kwargs - assert "addMemoryStrategies" in kwargs["memoryStrategies"] - assert kwargs["memoryStrategies"]["addMemoryStrategies"][0] == strategy + args, kwargs = mock_client.update_memory.call_args + assert "memoryStrategies" in kwargs + assert "addMemoryStrategies" in kwargs["memoryStrategies"] + added = kwargs["memoryStrategies"]["addMemoryStrategies"][0] + assert "semanticMemoryStrategy" in added + assert added["semanticMemoryStrategy"]["name"] == "Test Strategy" def test_add_strategy_wait_for_active(): """Test add_strategy with wait_for_active=True.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + client, mock_client = _create_client() - # Mock update_memory response (add_strategy uses update_memory internally) - mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} + # First: update_memory response + mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} - # Mock get_memory response to find the newly added strategy - mock_client.get_memory.return_value = { + # Subsequent: get_memory calls for strategy lookup and wait + mock_client.get_memory.side_effect = [ + # get_memory call to find newly added strategy + { "memory": { "id": "mem-123", - "strategies": [{"strategyId": "strat-new-123", "name": "Test Active Strategy", "status": "CREATING"}], + "status": "ACTIVE", + "strategies": [ + {"strategyId": "strat-new-123", "name": "Test Active Strategy", "status": "CREATING"}, + ], } - } - - # Test strategy addition with wait_for_active=True - strategy = {"semanticMemoryStrategy": {"name": "Test Active Strategy"}} + }, + # _wait_for_strategy_active polling: strategy now ACTIVE + { + "memory": { + "id": "mem-123", + "status": "ACTIVE", + "strategies": [ + {"strategyId": "strat-new-123", "name": "Test Active Strategy", "status": "ACTIVE"}, + ], + } + }, + ] - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - # Mock the _wait_for_strategy_active method to avoid actual waiting - with patch.object(client, "_wait_for_strategy_active") as mock_wait: - mock_wait.return_value = {"id": "mem-123", "status": "ACTIVE"} + strategy = {"semanticMemoryStrategy": {"name": "Test Active Strategy"}} + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + with patch("time.time", return_value=0): + with patch("time.sleep"): result = client.add_strategy("mem-123", strategy, wait_for_active=True, max_wait=120, poll_interval=5) assert result["id"] == "mem-123" assert mock_client.update_memory.called - # Verify strategy was passed correctly to update_memory - args, kwargs = mock_client.update_memory.call_args - assert "memoryStrategies" in kwargs - assert "addMemoryStrategies" in kwargs["memoryStrategies"] - assert kwargs["memoryStrategies"]["addMemoryStrategies"][0] == strategy - - # Verify get_memory was called to find the newly added strategy - assert mock_client.get_memory.called - get_args, get_kwargs = mock_client.get_memory.call_args - assert get_kwargs["memoryId"] == "mem-123" - - # Verify _wait_for_strategy_active was called with correct parameters - mock_wait.assert_called_once_with("mem-123", "strat-new-123", 120, 5) - def test_get_strategy(): """Test get_strategy functionality.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # Mock get_memory response with strategies - mock_client.get_memory.return_value = { - "memory": { - "id": "mem-123", - "strategies": [ - {"strategyId": "strat-1", "name": "Strategy 1", "type": "SEMANTIC"}, - {"strategyId": "strat-2", "name": "Strategy 2", "type": "SUMMARY"}, - ], - } + client, mock_client = _create_client() + + mock_client.get_memory.return_value = { + "memory": { + "id": "mem-123", + "strategies": [ + {"strategyId": "strat-1", "name": "Strategy 1", "type": "SEMANTIC"}, + {"strategyId": "strat-2", "name": "Strategy 2", "type": "SUMMARY"}, + ], } + } - # Test getting specific strategy - result = client.get_strategy("mem-123", "strat-1") + result = client.get_strategy("mem-123", "strat-1") - assert result["strategyId"] == "strat-1" - assert result["name"] == "Strategy 1" - assert result["type"] == "SEMANTIC" + assert result["strategyId"] == "strat-1" + assert result["name"] == "Strategy 1" def test_update_strategy(): """Test update_strategy functionality.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # Mock update_memory response (update_strategy uses update_memory internally) - mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} - - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - # Test strategy update - result = client.update_strategy( - memory_id="mem-123", - strategy_id="strat-456", - description="Updated strategy description", - namespaces=["custom/namespace1/", "custom/namespace2/"], - configuration={"modelId": "test-model"}, - ) + mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} - assert result["id"] == "mem-123" - assert mock_client.update_memory.called + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + result = client.update_strategy( + memory_id="mem-123", + strategy_id="strat-456", + description="Updated strategy description", + namespaces=["custom/namespace1/", "custom/namespace2/"], + ) - # Verify correct parameters were passed - args, kwargs = mock_client.update_memory.call_args - assert kwargs["memoryId"] == "mem-123" - assert "memoryStrategies" in kwargs - assert "modifyMemoryStrategies" in kwargs["memoryStrategies"] + assert result["id"] == "mem-123" + assert mock_client.update_memory.called - # Verify the strategy modification details - modify_strategy = kwargs["memoryStrategies"]["modifyMemoryStrategies"][0] - assert modify_strategy["memoryStrategyId"] == "strat-456" - assert modify_strategy["description"] == "Updated strategy description" - assert modify_strategy["namespaces"] == ["custom/namespace1/", "custom/namespace2/"] - assert modify_strategy["configuration"] == {"modelId": "test-model"} + args, kwargs = mock_client.update_memory.call_args + assert kwargs["memoryId"] == "mem-123" + assert "memoryStrategies" in kwargs + assert "modifyMemoryStrategies" in kwargs["memoryStrategies"] + modify_strategy = kwargs["memoryStrategies"]["modifyMemoryStrategies"][0] + assert modify_strategy["memoryStrategyId"] == "strat-456" + assert modify_strategy["description"] == "Updated strategy description" + assert modify_strategy["namespaces"] == ["custom/namespace1/", "custom/namespace2/"] -def test_error_handling(): - """Test error handling.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - # Mock the client to raise an error - mock_client = MagicMock() - client.client = mock_client +def test_remove_strategy(): + """Test remove_strategy functionality.""" + client, mock_client = _create_client() - error_response = {"Error": {"Code": "ValidationException", "Message": "Invalid parameter"}} - mock_client.create_memory.side_effect = ClientError(error_response, "CreateMemory") + mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} - try: - client.create_memory(name="Test Memory") - raise AssertionError("Error was not raised as expected") - except ClientError as e: - assert "ValidationException" in str(e) + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + result = client.remove_strategy(memory_id="mem-123", strategy_id="strat-456") + assert result["id"] == "mem-123" + assert mock_client.update_memory.called -def test_wait_for_strategy_active(): - """Test _wait_for_strategy_active helper method.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + args, kwargs = mock_client.update_memory.call_args + assert kwargs["memoryId"] == "mem-123" + assert "memoryStrategies" in kwargs + assert "deleteMemoryStrategies" in kwargs["memoryStrategies"] + deleted = kwargs["memoryStrategies"]["deleteMemoryStrategies"] + assert deleted == [{"memoryStrategyId": "strat-456"}] - # Mock get_memory response - strategy becomes ACTIVE - mock_client.get_memory.return_value = { - "memory": { - "id": "mem-123", - "strategies": [{"strategyId": "strat-456", "status": "ACTIVE", "name": "Test Strategy"}], - } - } - with patch("time.time", return_value=0): - with patch("time.sleep"): - # Test _wait_for_strategy_active - result = client._wait_for_strategy_active("mem-123", "strat-456", max_wait=60, poll_interval=5) +def test_error_handling(): + """Test error handling.""" + client, mock_client = _create_client() - assert result["id"] == "mem-123" - assert mock_client.get_memory.called + error_response = {"Error": {"Code": "ValidationException", "Message": "Invalid parameter"}} + mock_client.create_memory.side_effect = ClientError(error_response, "CreateMemory") - # Verify correct parameters - args, kwargs = mock_client.get_memory.call_args - assert kwargs["memoryId"] == "mem-123" + try: + client.create_memory(name="Test Memory") + raise AssertionError("Error was not raised as expected") + except ClientError as e: + assert "ValidationException" in str(e) def test_create_memory_with_strategies(): """Test create_memory with memory strategies.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + mock_client.create_memory.return_value = { + "memory": {"id": "mem-456", "name": "Memory with Strategies", "status": "CREATING"} + } - # Mock successful response - mock_client.create_memory.return_value = { - "memory": {"id": "mem-456", "name": "Memory with Strategies", "status": "CREATING"} - } + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + strategies = [{"semanticMemoryStrategy": {"name": "Strategy 1"}}] + result = client.create_memory( + name="Memory with Strategies", + description="Test with strategies", + strategies=strategies, + event_expiry_days=120, + memory_execution_role_arn="arn:aws:iam::123456789012:role/MemoryRole", + ) - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - # Test memory creation with strategies - strategies = [{"semanticMemoryStrategy": {"name": "Strategy 1"}}] - result = client.create_memory( - name="Memory with Strategies", - description="Test with strategies", - strategies=strategies, - event_expiry_days=120, - memory_execution_role_arn="arn:aws:iam::123456789012:role/MemoryRole", - ) - - assert result["id"] == "mem-456" - assert mock_client.create_memory.called + assert result["id"] == "mem-456" + assert mock_client.create_memory.called - # Verify all parameters were passed - args, kwargs = mock_client.create_memory.call_args - assert kwargs["name"] == "Memory with Strategies" - assert kwargs["description"] == "Test with strategies" - assert kwargs["memoryStrategies"] == strategies - assert kwargs["eventExpiryDuration"] == 120 - assert kwargs["memoryExecutionRoleArn"] == "arn:aws:iam::123456789012:role/MemoryRole" + args, kwargs = mock_client.create_memory.call_args + assert kwargs["name"] == "Memory with Strategies" + assert kwargs["description"] == "Test with strategies" + assert kwargs["eventExpiryDuration"] == 120 + assert kwargs["memoryExecutionRoleArn"] == "arn:aws:iam::123456789012:role/MemoryRole" + assert kwargs["memoryStrategies"] == strategies def test_list_memories_with_pagination(): """Test list_memories with pagination.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + client, mock_client = _create_client() - # Mock paginated responses - first_batch = [{"id": f"mem-{i}", "name": f"Memory {i}", "status": "ACTIVE"} for i in range(1, 101)] - second_batch = [{"id": f"mem-{i}", "name": f"Memory {i}", "status": "ACTIVE"} for i in range(101, 151)] + first_batch = [{"id": f"mem-{i}", "name": f"Memory {i}", "status": "ACTIVE"} for i in range(1, 101)] + second_batch = [{"id": f"mem-{i}", "name": f"Memory {i}", "status": "ACTIVE"} for i in range(101, 151)] - mock_client.list_memories.side_effect = [ - {"memories": first_batch, "nextToken": "token-123"}, - {"memories": second_batch, "nextToken": None}, - ] + mock_client.list_memories.side_effect = [ + {"memories": first_batch, "nextToken": "token-123"}, + {"memories": second_batch, "nextToken": None}, + ] - # Test with max_results requiring pagination - result = client.list_memories(max_results=150) + result = client.list_memories(max_results=150) - assert len(result) == 150 - assert result[0]["id"] == "mem-1" - assert result[149]["id"] == "mem-150" + assert len(result) == 150 + assert result[0]["id"] == "mem-1" + assert result[149]["id"] == "mem-150" + assert result[0]["strategyCount"] == 0 - # Verify two API calls were made - assert mock_client.list_memories.call_count == 2 + assert mock_client.list_memories.call_count == 2 def test_update_memory_minimal(): """Test update_memory with minimal parameters.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # Mock response - mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "ACTIVE"}} - - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - # Test minimal update (only memory_id) - result = client.update_memory(memory_id="mem-123") - - assert result["id"] == "mem-123" - assert mock_client.update_memory.called - - # Verify minimal parameters - args, kwargs = mock_client.update_memory.call_args - assert kwargs["memoryId"] == "mem-123" - assert kwargs["clientToken"] == "12345678-1234-5678-1234-567812345678" - + client, mock_client = _create_client() -def test_wait_for_status_timeout(): - """Test _wait_for_status with timeout.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "ACTIVE"}} - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + result = client.update_memory(memory_id="mem-123") - # Mock get_memory to always return CREATING (never becomes ACTIVE) - mock_client.get_memory.return_value = {"memory": {"id": "mem-timeout", "status": "CREATING", "strategies": []}} - - # Mock time to simulate timeout - provide enough values for all calls - time_values = [0] + [i * 10 for i in range(1, 35)] + [301] # Enough values for multiple checks - with patch("time.time", side_effect=time_values): - with patch("time.sleep"): - try: - client._wait_for_status( - memory_id="mem-timeout", target_status="ACTIVE", max_wait=300, poll_interval=10 - ) - raise AssertionError("TimeoutError was not raised") - except TimeoutError as e: - assert "did not reach status ACTIVE within 300 seconds" in str(e) - - -def test_wait_for_status_failure(): - """Test _wait_for_status with FAILED status.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # Mock get_memory to return FAILED status - mock_client.get_memory.return_value = { - "memory": {"id": "mem-failed", "status": "FAILED", "failureReason": "Configuration error", "strategies": []} - } - - with patch("time.time", return_value=0): - with patch("time.sleep"): - try: - client._wait_for_status( - memory_id="mem-failed", target_status="ACTIVE", max_wait=300, poll_interval=10 - ) - raise AssertionError("RuntimeError was not raised") - except RuntimeError as e: - assert "Memory operation failed: Configuration error" in str(e) - - -def test_wait_for_strategy_active_timeout(): - """Test _wait_for_strategy_active with timeout.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # Mock get_memory response - strategy never becomes ACTIVE - mock_client.get_memory.return_value = { - "memory": {"id": "mem-123", "strategies": [{"strategyId": "strat-timeout", "status": "CREATING"}]} - } - - # Mock time to simulate timeout - provide enough values for multiple calls - time_values = [0] + [i * 10 for i in range(1, 35)] + [301] - with patch("time.time", side_effect=time_values): - with patch("time.sleep"): - try: - client._wait_for_strategy_active("mem-123", "strat-timeout", max_wait=300, poll_interval=10) - raise AssertionError("TimeoutError was not raised") - except TimeoutError as e: - assert "Strategy strat-timeout did not become ACTIVE within 300 seconds" in str(e) - - -def test_wait_for_strategy_active_not_found(): - """Test _wait_for_strategy_active when strategy is not found.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # Mock get_memory response - strategy doesn't exist - mock_client.get_memory.return_value = { - "memory": {"id": "mem-123", "strategies": [{"strategyId": "strat-other", "status": "ACTIVE"}]} - } + assert result["id"] == "mem-123" + assert mock_client.update_memory.called - # Mock time to simulate timeout - provide enough values for multiple calls - time_values = [0] + [i * 5 for i in range(1, 15)] + [61] - with patch("time.time", side_effect=time_values): - with patch("time.sleep"): - try: - client._wait_for_strategy_active("mem-123", "strat-nonexistent", max_wait=60, poll_interval=5) - raise AssertionError("TimeoutError was not raised") - except TimeoutError as e: - assert "Strategy strat-nonexistent did not become ACTIVE within 60 seconds" in str(e) + args, kwargs = mock_client.update_memory.call_args + assert kwargs["memoryId"] == "mem-123" + assert kwargs["clientToken"] == "12345678-1234-5678-1234-567812345678" def test_get_strategy_not_found(): """Test get_strategy when strategy doesn't exist.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + client, mock_client = _create_client() - # Mock get_memory response without the requested strategy - mock_client.get_memory.return_value = { - "memory": { - "id": "mem-123", - "strategies": [{"strategyId": "strat-other", "name": "Other Strategy", "type": "SEMANTIC"}], - } + mock_client.get_memory.return_value = { + "memory": { + "id": "mem-123", + "strategies": [{"strategyId": "strat-other", "name": "Other Strategy", "type": "SEMANTIC"}], } + } - try: - client.get_strategy("mem-123", "strat-nonexistent") - raise AssertionError("ValueError was not raised") - except ValueError as e: - assert "Strategy strat-nonexistent not found in memory mem-123" in str(e) + try: + client.get_strategy("mem-123", "strat-nonexistent") + raise AssertionError("ValueError was not raised") + except ValueError as e: + assert "Strategy strat-nonexistent not found in memory mem-123" in str(e) def test_delete_memory_wait_for_deletion_timeout(): """Test delete_memory with wait_for_deletion timeout.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + client, mock_client = _create_client() - # Mock delete_memory response - mock_client.delete_memory.return_value = {"status": "DELETING"} + mock_client.delete_memory.return_value = {"status": "DELETING"} - # Mock get_memory to always succeed (memory never gets deleted) - mock_client.get_memory.return_value = {"memory": {"id": "mem-persistent", "status": "DELETING"}} + # Mock get_memory to always succeed (memory never gets deleted) + mock_client.get_memory.return_value = {"memory": {"id": "mem-persistent", "status": "DELETING"}} - # Mock time to simulate timeout - # Provide enough values for multiple time.time() calls in the loop - with patch("time.time", side_effect=[0, 0, 0, 301, 301, 301]): - with patch("time.sleep"): - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - try: - client.delete_memory("mem-persistent", wait_for_deletion=True, max_wait=300, poll_interval=10) - raise AssertionError("TimeoutError was not raised") - except TimeoutError as e: - assert "Memory mem-persistent was not deleted within 300 seconds" in str(e) + with patch("time.time", side_effect=[0, 0, 0, 301, 301, 301]): + with patch("time.sleep"): + with patch( + "uuid.uuid4", + return_value=uuid.UUID("12345678-1234-5678-1234-567812345678"), + ): + try: + client.delete_memory( + "mem-persistent", + wait_for_deletion=True, + max_wait=300, + poll_interval=10, + ) + raise AssertionError("TimeoutError was not raised") + except TimeoutError as e: + assert "was not deleted within 300 seconds" in str(e) -def test_wait_for_status_with_strategy_check(): - """Test _wait_for_status with check_strategies=True and transitional strategies.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() +def test_wait_for_status_timeout(): + """Test _wait_for_status with timeout.""" + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # Mock get_memory responses - first with transitional strategy, then all active - mock_client.get_memory.side_effect = [ - { - "memory": { - "id": "mem-123", - "status": "ACTIVE", - "strategies": [ - {"strategyId": "strat-1", "status": "CREATING"}, # Transitional - {"strategyId": "strat-2", "status": "ACTIVE"}, # Already active - ], - } - }, - { - "memory": { - "id": "mem-123", - "status": "ACTIVE", - "strategies": [ - {"strategyId": "strat-1", "status": "ACTIVE"}, # Now active - {"strategyId": "strat-2", "status": "ACTIVE"}, - ], - } - }, - ] + mock_client.get_memory.return_value = {"memory": {"id": "mem-timeout", "status": "CREATING", "strategies": []}} - with patch("time.time", return_value=0): - with patch("time.sleep"): - # Test _wait_for_status with check_strategies=True - result = client._wait_for_status( - memory_id="mem-123", target_status="ACTIVE", max_wait=120, poll_interval=10, check_strategies=True + time_values = [0] + [i * 10 for i in range(1, 35)] + [301] + with patch("time.time", side_effect=time_values): + with patch("time.sleep"): + try: + client._wait_for_status( + "mem-timeout", + target_status=MemoryStatus.ACTIVE.value, + max_wait=300, + poll_interval=10, ) + raise AssertionError("TimeoutError was not raised") + except TimeoutError as e: + assert "did not reach status ACTIVE within 300 seconds" in str(e) - assert result["id"] == "mem-123" - assert result["status"] == "ACTIVE" - - # Should have made two calls - one found transitional strategy, second found all active - assert mock_client.get_memory.call_count == 2 +def test_wait_for_status_failure(): + """Test _wait_for_status with FAILED status.""" + client, mock_client = _create_client() + + mock_client.get_memory.return_value = { + "memory": { + "id": "mem-failed", + "status": "FAILED", + "failureReason": "Configuration error", + "strategies": [], + } + } -def test_add_strategy_strategy_not_found(): - """Test add_strategy when newly added strategy cannot be found.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + with patch("time.time", return_value=0): + with patch("time.sleep"): + try: + client._wait_for_status( + "mem-failed", + target_status=MemoryStatus.ACTIVE.value, + max_wait=300, + poll_interval=10, + ) + raise AssertionError("RuntimeError was not raised") + except RuntimeError as e: + assert "Memory operation failed: Configuration error" in str(e) - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - # Mock update_memory response - mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} +def test_wait_for_status_with_strategy_check(): + """Test _wait_for_status with transitional strategies.""" + client, mock_client = _create_client() - # Mock get_memory response without the newly added strategy - mock_client.get_memory.return_value = { + mock_client.get_memory.side_effect = [ + { "memory": { "id": "mem-123", "status": "ACTIVE", - "strategies": [], # No strategies found - } - } - - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - strategy = {"semanticMemoryStrategy": {"name": "Missing Strategy"}} - - # The actual implementation just logs a warning and returns the memory - # It doesn't raise an exception - result = client.add_strategy("mem-123", strategy, wait_for_active=True) - - # Should return the memory object from get_memory (since wait_for_active=True) - assert result["id"] == "mem-123" - assert result["status"] == "ACTIVE" - - -def test_initialization_with_env_vars(): - """Test initialization with environment variables.""" - with patch("boto3.client") as mock_boto_client: - with patch("os.getenv") as mock_getenv: - # Mock environment variables - use the correct names from controlplane.py - env_vars = { - "BEDROCK_AGENTCORE_CONTROL_ENDPOINT": "https://custom-control.amazonaws.com", - "BEDROCK_AGENTCORE_CONTROL_SERVICE": "custom-control-service", + "strategies": [ + {"strategyId": "strat-1", "status": "CREATING"}, + {"strategyId": "strat-2", "status": "ACTIVE"}, + ], } - mock_getenv.side_effect = lambda key, default=None: env_vars.get(key, default) - - # Test initialization with custom environment - MemoryControlPlaneClient() - - # Verify boto3.client was called with custom endpoint - mock_boto_client.assert_called_with( - "custom-control-service", region_name="us-west-2", endpoint_url="https://custom-control.amazonaws.com" - ) - - -def test_wait_for_status(): - """Test _wait_for_status helper method.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # Mock get_memory response - memory becomes ACTIVE - mock_client.get_memory.return_value = { + }, + { "memory": { "id": "mem-123", "status": "ACTIVE", @@ -795,165 +565,128 @@ def test_wait_for_status(): {"strategyId": "strat-2", "status": "ACTIVE"}, ], } - } - - with patch("time.time", return_value=0): - with patch("time.sleep"): - # Test _wait_for_status with check_strategies=True - result = client._wait_for_status( - memory_id="mem-123", target_status="ACTIVE", max_wait=120, poll_interval=10, check_strategies=True - ) - - assert result["id"] == "mem-123" - assert result["status"] == "ACTIVE" - assert mock_client.get_memory.called + }, + ] + + with patch("time.time", return_value=0): + with patch("time.sleep"): + result = client._wait_for_status( + "mem-123", + target_status=MemoryStatus.ACTIVE.value, + max_wait=120, + poll_interval=10, + ) - # Verify correct parameters - args, kwargs = mock_client.get_memory.call_args - assert kwargs["memoryId"] == "mem-123" + assert result["id"] == "mem-123" + assert result["status"] == "ACTIVE" + assert mock_client.get_memory.call_count == 2 def test_get_memory_client_error(): """Test get_memory with ClientError.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + error_response = {"Error": {"Code": "ResourceNotFoundException", "Message": "Memory not found"}} + mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") - # Mock ClientError - error_response = {"Error": {"Code": "ResourceNotFoundException", "Message": "Memory not found"}} - mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") - - try: - client.get_memory("nonexistent-mem-123") - raise AssertionError("ClientError was not raised") - except ClientError as e: - assert "ResourceNotFoundException" in str(e) + try: + client.get_memory("nonexistent-mem-123") + raise AssertionError("ClientError was not raised") + except ClientError as e: + assert "ResourceNotFoundException" in str(e) def test_list_memories_client_error(): """Test list_memories with ClientError.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + error_response = {"Error": {"Code": "AccessDeniedException", "Message": "Insufficient permissions"}} + mock_client.list_memories.side_effect = ClientError(error_response, "ListMemories") - # Mock ClientError - error_response = {"Error": {"Code": "AccessDeniedException", "Message": "Insufficient permissions"}} - mock_client.list_memories.side_effect = ClientError(error_response, "ListMemories") - - try: - client.list_memories(max_results=50) - raise AssertionError("ClientError was not raised") - except ClientError as e: - assert "AccessDeniedException" in str(e) + try: + client.list_memories(max_results=50) + raise AssertionError("ClientError was not raised") + except ClientError as e: + assert "AccessDeniedException" in str(e) def test_update_memory_client_error(): """Test update_memory with ClientError.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + error_response = {"Error": {"Code": "ValidationException", "Message": "Invalid memory parameters"}} + mock_client.update_memory.side_effect = ClientError(error_response, "UpdateMemory") - # Mock ClientError - error_response = {"Error": {"Code": "ValidationException", "Message": "Invalid memory parameters"}} - mock_client.update_memory.side_effect = ClientError(error_response, "UpdateMemory") - - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - try: - client.update_memory(memory_id="mem-123", description="Updated description") - raise AssertionError("ClientError was not raised") - except ClientError as e: - assert "ValidationException" in str(e) + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + try: + client.update_memory(memory_id="mem-123", description="Updated description") + raise AssertionError("ClientError was not raised") + except ClientError as e: + assert "ValidationException" in str(e) def test_delete_memory_client_error(): """Test delete_memory with ClientError.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + error_response = {"Error": {"Code": "ConflictException", "Message": "Memory is in use"}} + mock_client.delete_memory.side_effect = ClientError(error_response, "DeleteMemory") - # Mock ClientError - error_response = {"Error": {"Code": "ConflictException", "Message": "Memory is in use"}} - mock_client.delete_memory.side_effect = ClientError(error_response, "DeleteMemory") - - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - try: - client.delete_memory("mem-in-use") - raise AssertionError("ClientError was not raised") - except ClientError as e: - assert "ConflictException" in str(e) + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + try: + client.delete_memory("mem-in-use") + raise AssertionError("ClientError was not raised") + except ClientError as e: + assert "ConflictException" in str(e) def test_get_strategy_client_error(): """Test get_strategy with ClientError from get_memory.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + error_response = {"Error": {"Code": "ThrottlingException", "Message": "Request throttled"}} + mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") - # Mock ClientError from get_memory call - error_response = {"Error": {"Code": "ThrottlingException", "Message": "Request throttled"}} - mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") - - try: - client.get_strategy("mem-123", "strat-456") - raise AssertionError("ClientError was not raised") - except ClientError as e: - assert "ThrottlingException" in str(e) + try: + client.get_strategy("mem-123", "strat-456") + raise AssertionError("ClientError was not raised") + except ClientError as e: + assert "ThrottlingException" in str(e) -def test_wait_for_strategy_active_client_error(): - """Test _wait_for_strategy_active with ClientError.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() +def test_wait_for_status_client_error(): + """Test _wait_for_status with ClientError.""" + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + error_response = {"Error": {"Code": "InternalServerError", "Message": "Internal server error"}} + mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") - # Mock ClientError - error_response = {"Error": {"Code": "ServiceException", "Message": "Internal service error"}} - mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") + with patch("time.time", return_value=0): + with patch("time.sleep"): + try: + client._wait_for_status( + "mem-123", + target_status=MemoryStatus.ACTIVE.value, + max_wait=120, + poll_interval=10, + ) + raise AssertionError("ClientError was not raised") + except ClientError as e: + assert "InternalServerError" in str(e) - with patch("time.time", return_value=0): - with patch("time.sleep"): - try: - client._wait_for_strategy_active("mem-123", "strat-456", max_wait=60, poll_interval=5) - raise AssertionError("ClientError was not raised") - except ClientError as e: - assert "ServiceException" in str(e) +def test_constructor_client_param(): + """Test that passing client= in constructor sets self.client directly.""" + mock_client = MagicMock() + client = MemoryControlPlaneClient(client=mock_client) -def test_wait_for_status_client_error(): - """Test _wait_for_status with ClientError.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + assert client.client is mock_client - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - # Mock ClientError - error_response = {"Error": {"Code": "InternalServerError", "Message": "Internal server error"}} - mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") +def test_constructor_default_client(): + """Test that omitting client= in constructor creates a default boto3 client.""" + with patch("boto3.client") as mock_boto: + mock_boto.return_value = MagicMock() + client = MemoryControlPlaneClient() - with patch("time.time", return_value=0): - with patch("time.sleep"): - try: - client._wait_for_status(memory_id="mem-123", target_status="ACTIVE", max_wait=120, poll_interval=10) - raise AssertionError("ClientError was not raised") - except ClientError as e: - assert "InternalServerError" in str(e) + assert client.client is mock_boto.return_value