From 2bfcf893ca65ec2ba1bccc3c385fec41818a3543 Mon Sep 17 00:00:00 2001 From: "T.J Ariyawansa" Date: Tue, 10 Feb 2026 15:02:11 -0500 Subject: [PATCH 1/4] Deprecate MemoryControlPlaneClient as pass-through to MemoryClient MemoryControlPlaneClient and MemoryClient have significant functional overlap. MemoryClient is a superset that provides all control plane operations plus data plane features, response normalization, and typed strategy helpers. Maintaining both creates a maintenance burden and user confusion. This refactors MemoryControlPlaneClient to be a thin pass-through that delegates all operations to an internal MemoryClient instance, and emits a DeprecationWarning on instantiation directing users to migrate. Closes #247 --- src/bedrock_agentcore/memory/controlplane.py | 374 ++--- .../memory/test_controlplane.py | 1208 +++++++---------- 2 files changed, 591 insertions(+), 991 deletions(-) diff --git a/src/bedrock_agentcore/memory/controlplane.py b/src/bedrock_agentcore/memory/controlplane.py index 25e75d09..333c23c7 100644 --- a/src/bedrock_agentcore/memory/controlplane.py +++ b/src/bedrock_agentcore/memory/controlplane.py @@ -1,46 +1,60 @@ -"""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 .client import MemoryClient +from .constants import MemoryStatus logger = logging.getLogger(__name__) 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"): """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) """ + 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. See: https://github.com/aws/bedrock-agentcore-sdk-python/issues/247", + DeprecationWarning, + stacklevel=2, + ) self.region_name = region_name self.environment = environment + self._memory_client = MemoryClient(region_name=region_name) - self.endpoint = os.getenv( - "BEDROCK_AGENTCORE_CONTROL_ENDPOINT", f"https://bedrock-agentcore-control.{region_name}.amazonaws.com" - ) + logger.info("Initialized MemoryControlPlaneClient for %s in %s", environment, region_name) - 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) + @property + def client(self): + """The underlying boto3 control plane client.""" + return self._memory_client.gmcp_client - logger.info("Initialized MemoryControlPlaneClient for %s in %s", environment, region_name) + @client.setter + def client(self, value): + """Allow overriding the underlying boto3 client (used by tests).""" + self._memory_client.gmcp_client = value # ==================== MEMORY OPERATIONS ==================== @@ -70,36 +84,23 @@ def create_memory( Returns: Created memory object """ - params = { - "name": name, - "eventExpiryDuration": event_expiry_days, - "clientToken": str(uuid.uuid4()), - } - - if description: - params["description"] = description - - if memory_execution_role_arn: - params["memoryExecutionRoleArn"] = memory_execution_role_arn - - if strategies: - params["memoryStrategies"] = strategies - - try: - response = self.client.create_memory(**params) - memory = response["memory"] - memory_id = memory["id"] - - logger.info("Created memory: %s", memory_id) - - if wait_for_active: - return self._wait_for_memory_active(memory_id, max_wait, poll_interval) - - return memory - - except ClientError as e: - logger.error("Failed to create memory: %s", e) - raise + if wait_for_active: + return self._memory_client.create_memory_and_wait( + name=name, + strategies=strategies or [], + description=description, + event_expiry_days=event_expiry_days, + memory_execution_role_arn=memory_execution_role_arn, + max_wait=max_wait, + poll_interval=poll_interval, + ) + return self._memory_client.create_memory( + name=name, + strategies=strategies or [], + description=description, + event_expiry_days=event_expiry_days, + memory_execution_role_arn=memory_execution_role_arn, + ) def get_memory(self, memory_id: str, include_strategies: bool = True) -> Dict[str, Any]: """Get a memory resource by ID. @@ -138,32 +139,7 @@ def list_memories(self, max_results: int = 100) -> List[Dict[str, Any]]: Returns: List of memory summaries """ - try: - memories = [] - next_token = None - - while len(memories) < max_results: - params = {"maxResults": min(100, max_results - len(memories))} - if next_token: - params["nextToken"] = next_token - - response = self.client.list_memories(**params) - batch = response.get("memories", []) - memories.extend(batch) - - next_token = response.get("nextToken") - if not next_token or len(memories) >= max_results: - break - - # Add strategy count to each memory summary - for memory in memories: - memory["strategyCount"] = 0 # List memories doesn't include strategies - - return memories[:max_results] - - except ClientError as e: - logger.error("Failed to list memories: %s", e) - raise + return self._memory_client.list_memories(max_results=max_results) def update_memory( self, @@ -245,7 +221,7 @@ 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]: @@ -279,37 +255,28 @@ def delete_memory( logger.info( "Waiting for %d strategies to become ACTIVE before deletion", len(transitional_strategies) ) - self._wait_for_status( - memory_id=memory_id, - target_status=MemoryStatus.ACTIVE.value, - max_wait=max_wait, - poll_interval=poll_interval, - check_strategies=True, - ) + start_time = time.time() + while time.time() - start_time < max_wait: + memory = self.get_memory(memory_id) + strategies = memory.get("strategies", []) + all_ready = all( + s.get("status") in [MemoryStatus.ACTIVE.value, MemoryStatus.FAILED.value] + for s in strategies + ) + if all_ready: + break + time.sleep(poll_interval) except Exception as e: logger.warning("Error waiting for strategies to become ACTIVE: %s", e) - # Now delete the memory - response = self.client.delete_memory(memoryId=memory_id, clientToken=str(uuid.uuid4())) + if wait_for_deletion: + return self._memory_client.delete_memory_and_wait( + memory_id=memory_id, + max_wait=max_wait, + poll_interval=poll_interval, + ) - logger.info("Initiated deletion of memory: %s", memory_id) - - if not wait_for_deletion: - return response - - # Wait for deletion to complete - start_time = time.time() - while time.time() - start_time < max_wait: - try: - self.client.get_memory(memoryId=memory_id) - time.sleep(poll_interval) - except ClientError as e: - if e.response["Error"]["Code"] == "ResourceNotFoundException": - logger.info("Memory %s successfully deleted", memory_id) - return response - raise - - raise TimeoutError(f"Memory {memory_id} was not deleted within {max_wait} seconds") + return self._memory_client.delete_memory(memory_id=memory_id) except ClientError as e: logger.error("Failed to delete memory: %s", e) @@ -335,44 +302,20 @@ def add_strategy( poll_interval: Seconds between status checks if wait_for_active is True Returns: - Updated memory object with strategyId field + Updated memory object """ - # Get the strategy type and name for identification - strategy_type = list(strategy.keys())[0] # e.g., 'semanticMemoryStrategy' - strategy_name = strategy[strategy_type].get("name") - - logger.info("Adding strategy %s of type %s to memory %s", strategy_name, strategy_type, memory_id) - - # Use update_memory with add_strategies parameter but don't wait for memory - memory = self.update_memory( + if wait_for_active: + return self._memory_client.update_memory_strategies_and_wait( + memory_id=memory_id, + add_strategies=[strategy], + max_wait=max_wait, + poll_interval=poll_interval, + ) + return self._memory_client.update_memory_strategies( memory_id=memory_id, add_strategies=[strategy], - wait_for_active=False, # Don't wait for memory, we'll check strategy specifically ) - # If we need to wait for the strategy to become active - if wait_for_active: - # First, get the memory again to ensure we have the latest state - memory = self.get_memory(memory_id) - - # Find the newly added strategy by matching name - strategies = memory.get("strategies", []) - strategy_id = None - - for s in strategies: - # Match by name since that's unique within a memory - if s.get("name") == strategy_name: - strategy_id = s.get("strategyId") - logger.info("Found newly added strategy %s with ID %s", strategy_name, strategy_id) - break - - if strategy_id: - return self._wait_for_strategy_active(memory_id, strategy_id, max_wait, poll_interval) - else: - logger.warning("Could not identify newly added strategy %s to wait for activation", strategy_name) - - return memory - def get_strategy(self, memory_id: str, strategy_id: str) -> Dict[str, Any]: """Get a specific strategy from a memory resource. @@ -384,8 +327,7 @@ def get_strategy(self, memory_id: str, strategy_id: str) -> Dict[str, Any]: Strategy details """ try: - memory = self.get_memory(memory_id) - strategies = memory.get("strategies", []) + strategies = self._memory_client.get_memory_strategies(memory_id) for strategy in strategies: if strategy.get("strategyId") == strategy_id: @@ -423,7 +365,6 @@ def update_strategy( Returns: Updated memory object """ - # Note: API expects memoryStrategyId for input but returns strategyId in response modify_config: Dict = {"memoryStrategyId": strategy_id} if description is not None: @@ -435,19 +376,18 @@ def update_strategy( if configuration is not None: modify_config["configuration"] = configuration - # Use update_memory with modify_strategies parameter but don't wait for memory - memory = self.update_memory( + if wait_for_active: + return self._memory_client.update_memory_strategies_and_wait( + memory_id=memory_id, + modify_strategies=[modify_config], + max_wait=max_wait, + poll_interval=poll_interval, + ) + return self._memory_client.update_memory_strategies( memory_id=memory_id, modify_strategies=[modify_config], - wait_for_active=False, # Don't wait for memory, we'll check strategy specifically ) - # If we need to wait for the strategy to become active - if wait_for_active: - return self._wait_for_strategy_active(memory_id, strategy_id, max_wait, poll_interval) - - return memory - def remove_strategy( self, memory_id: str, @@ -468,146 +408,56 @@ def remove_strategy( Returns: Updated memory object """ - # For remove_strategy, we only need to wait for memory to be active - # since the strategy will be gone - return self.update_memory( + if wait_for_active: + return self._memory_client.update_memory_strategies_and_wait( + memory_id=memory_id, + delete_strategy_ids=[strategy_id], + max_wait=max_wait, + poll_interval=poll_interval, + ) + return self._memory_client.update_memory_strategies( memory_id=memory_id, delete_strategy_ids=[strategy_id], - wait_for_active=wait_for_active, - max_wait=max_wait, - poll_interval=poll_interval, ) # ==================== HELPER METHODS ==================== def _wait_for_memory_active(self, memory_id: str, max_wait: int, poll_interval: int) -> Dict[str, Any]: - """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 - ) - - def _wait_for_strategy_active( - self, memory_id: str, strategy_id: str, max_wait: int, poll_interval: int - ) -> Dict[str, Any]: - """Wait for specific memory strategy to become ACTIVE.""" - logger.info("Waiting for strategy %s to become ACTIVE (max wait: %d seconds)...", strategy_id, max_wait) - - start_time = time.time() - last_status = None - - while time.time() - start_time < max_wait: - try: - memory = self.get_memory(memory_id) - strategies = memory.get("strategies", []) - - for strategy in strategies: - if strategy.get("strategyId") == strategy_id: - status = strategy["status"] - - # Log status changes - if status != last_status: - logger.info("Strategy %s status: %s", strategy_id, status) - last_status = status - - if status == MemoryStatus.ACTIVE.value: - elapsed = time.time() - start_time - logger.info("Strategy %s is now ACTIVE (took %.1f seconds)", strategy_id, elapsed) - return memory - elif status == MemoryStatus.FAILED.value: - failure_reason = strategy.get("failureReason", "Unknown") - raise RuntimeError(f"Strategy {strategy_id} failed to activate: {failure_reason}") - - break - else: - logger.warning("Strategy %s not found in memory %s", strategy_id, memory_id) - - # Wait before checking again - time.sleep(poll_interval) - - except ClientError as e: - logger.error("Error checking strategy status: %s", e) - raise - - elapsed = time.time() - start_time - raise TimeoutError( - f"Strategy {strategy_id} did not become ACTIVE within {max_wait} seconds (last status: {last_status})" - ) - - def _wait_for_status( - 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. - - Args: - memory_id: The ID of the memory to check - target_status: The status to wait for (e.g., "ACTIVE") - max_wait: Maximum time to wait in seconds - poll_interval: Time between status checks in seconds - check_strategies: Whether to also check that all strategies are in the target status + """Wait for memory and all strategies to reach ACTIVE state. - Returns: - The memory object once it reaches the target status - - Raises: - TimeoutError: If the memory doesn't reach the target status within max_wait - RuntimeError: If the memory or any strategy reaches a FAILED state + Used by update_memory(wait_for_active=True). """ - logger.info("Waiting for memory %s to reach status %s...", memory_id, target_status) + logger.info("Waiting for memory %s to become ACTIVE...", memory_id) start_time = time.time() last_memory_status = None - strategy_statuses = {} while time.time() - start_time < max_wait: try: memory = self.get_memory(memory_id) status = memory.get("status") - # Log status changes for memory if status != last_memory_status: logger.info("Memory %s status: %s", memory_id, status) last_memory_status = status - if status == target_status: - # Check if all strategies are also in the target status - if check_strategies and target_status == MemoryStatus.ACTIVE.value: - strategies = memory.get("strategies", []) - all_strategies_active = True - - for strategy in strategies: - strategy_id = strategy.get("strategyId") - strategy_status = strategy.get("status") - - # Log strategy status changes - if ( - strategy_id not in strategy_statuses - or strategy_statuses[strategy_id] != strategy_status - ): - logger.info("Strategy %s status: %s", strategy_id, strategy_status) - strategy_statuses[strategy_id] = strategy_status - - if strategy_status != target_status: - if strategy_status == MemoryStatus.FAILED.value: - failure_reason = strategy.get("failureReason", "Unknown") - raise RuntimeError(f"Strategy {strategy_id} failed: {failure_reason}") - - all_strategies_active = False - - if not all_strategies_active: - logger.info( - "Memory %s is %s but %d strategies are still processing", - memory_id, - target_status, - len([s for s in strategies if s.get("status") != target_status]), - ) - time.sleep(poll_interval) - continue + if status == MemoryStatus.ACTIVE.value: + strategies = memory.get("strategies", []) + all_strategies_active = all( + s.get("status") in [MemoryStatus.ACTIVE.value, MemoryStatus.FAILED.value] for s in strategies + ) + + if not all_strategies_active: + logger.info( + "Memory %s is ACTIVE but %d strategies are still processing", + memory_id, + len([s for s in strategies if s.get("status") != MemoryStatus.ACTIVE.value]), + ) + time.sleep(poll_interval) + continue elapsed = time.time() - start_time - logger.info( - "Memory %s and all strategies are now %s (took %.1f seconds)", memory_id, target_status, elapsed - ) + logger.info("Memory %s and all strategies are now ACTIVE (took %.1f seconds)", memory_id, elapsed) return memory elif status == MemoryStatus.FAILED.value: failure_reason = memory.get("failureReason", "Unknown") @@ -621,6 +471,6 @@ def _wait_for_status( elapsed = time.time() - start_time raise TimeoutError( - f"Memory {memory_id} did not reach status {target_status} within {max_wait} seconds " + f"Memory {memory_id} did not reach status ACTIVE within {max_wait} seconds " f"(elapsed: {elapsed:.1f}s)" ) diff --git a/tests/bedrock_agentcore/memory/test_controlplane.py b/tests/bedrock_agentcore/memory/test_controlplane.py index 38b691c0..7565e28c 100644 --- a/tests/bedrock_agentcore/memory/test_controlplane.py +++ b/tests/bedrock_agentcore/memory/test_controlplane.py @@ -1,305 +1,282 @@ """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 warning for all tests except the one that explicitly tests it +pytestmark = pytest.mark.filterwarnings("ignore::DeprecationWarning") -def test_create_memory(): - """Test create_memory functionality.""" + +def _create_client(): + """Helper to create a MemoryControlPlaneClient with mocked boto3 clients.""" with patch("boto3.client"): client = MemoryControlPlaneClient() + mock_client = MagicMock() + client.client = mock_client + return client, mock_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"} - } +def test_deprecation_warning(): + """Test that MemoryControlPlaneClient emits a deprecation warning.""" + with patch("boto3.client"): + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + MemoryControlPlaneClient() + assert len(w) == 1 + assert issubclass(w[0].category, DeprecationWarning) + assert "MemoryClient" in str(w[0].message) + assert "v1.4.0" in str(w[0].message) - 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") - assert result["id"] == "mem-123" - assert result["name"] == "Test Memory" - assert mock_client.create_memory.called +def test_create_memory(): + """Test create_memory functionality.""" + client, mock_client = _create_client() - # 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" + # 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 (delegates to MemoryClient.create_memory) + result = client.create_memory(name="Test Memory", description="Test description") -def test_get_memory(): - """Test get_memory functionality.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + assert result["id"] == "mem-123" + assert result["name"] == "Test Memory" + assert mock_client.create_memory.called - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + # 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" - # 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"}, - ], - } + +def test_get_memory(): + """Test get_memory functionality.""" + 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") + # Test get memory with strategies + 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() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + client, mock_client = _create_client() - # 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} + # 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} - # Test list memories - result = client.list_memories(max_results=50) + # Test list memories (delegates to MemoryClient.list_memories) + result = client.list_memories(max_results=50) - assert len(result) == 2 - assert result[0]["id"] == "mem-1" - assert result[0]["strategyCount"] == 0 # List doesn't include strategies + assert len(result) == 2 + assert result[0]["id"] == "mem-1" - # 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 response + 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")): + # Test memory update + 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" + # 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" def test_delete_memory(): """Test delete_memory functionality.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + client, mock_client = _create_client() - # Mock response - 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")): - # Test memory deletion - result = client.delete_memory("mem-123") + with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): + # Test memory deletion (delegates to MemoryClient.delete_memory) + 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" + # Verify correct parameters + 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() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + client, mock_client = _create_client() - # Mock get_memory response with strategies in transitional state - mock_client.get_memory.return_value = { + # Mock get_memory responses: first with transitional strategies, then all active + mock_client.get_memory.side_effect = [ + # First call from get_memory (initial check) + { "memory": { "id": "mem-123", "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"} + }, + # Second call from get_memory (polling - all active now) + { + "memory": { + "id": "mem-123", + "strategies": [ + {"strategyId": "strat-1", "status": "ACTIVE"}, + {"strategyId": "strat-2", "status": "ACTIVE"}, + ], + } + }, + ] - # Test memory deletion with wait_for_strategies=True - result = client.delete_memory("mem-123", wait_for_strategies=True) + # Mock delete_memory response + mock_client.delete_memory.return_value = {"status": "DELETING"} - assert result["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"): + result = client.delete_memory("mem-123", wait_for_strategies=True) - # Verify get_memory was called to check strategy status - assert mock_client.get_memory.called + assert result["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, - ) + # Verify get_memory was called to check strategy status + assert mock_client.get_memory.called - # Verify delete_memory was called - assert mock_client.delete_memory.called - args, kwargs = mock_client.delete_memory.call_args - assert kwargs["memoryId"] == "mem-123" + # Verify delete_memory was called + assert mock_client.delete_memory.called + args, kwargs = mock_client.delete_memory.call_args + assert kwargs["memoryId"] == "mem-123" 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 delete_memory response + 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") + # Mock get_memory to raise ResourceNotFoundException (memory is gone) + 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) + 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) - assert result["status"] == "DELETING" + 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" + # 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" - # 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" + # Verify get_memory was called (to check if memory is gone) + 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 update_memory response (add_strategy uses update_memory_strategies internally) + mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} - # Test strategy addition - strategy = {"semanticMemoryStrategy": {"name": "Test Strategy"}} + # Test strategy addition + 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 + # Verify strategy was passed in the update_memory call + 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"}} + # Mock update_memory response (via update_memory_strategies) + 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 = { - "memory": { - "id": "mem-123", - "strategies": [{"strategyId": "strat-new-123", "name": "Test Active Strategy", "status": "CREATING"}], - } + # Mock get_memory for _wait_for_memory_active (returns ACTIVE) + mock_client.get_memory.return_value = { + "memory": { + "id": "mem-123", + "status": "ACTIVE", + "strategies": [{"strategyId": "strat-new-123", "name": "Test Active Strategy", "status": "ACTIVE"}], } + } - # Test strategy addition with wait_for_active=True - strategy = {"semanticMemoryStrategy": {"name": "Test Active Strategy"}} - - 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"} + # Test strategy addition with wait_for_active=True + 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) + # Result comes from _wait_for_memory_active which returns normalized memory assert result["id"] == "mem-123" assert mock_client.update_memory.called @@ -307,486 +284,259 @@ def test_add_strategy_wait_for_active(): 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 get_memory response with strategies (used by get_memory_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"}, + ], } + } - # Test getting specific strategy - result = client.get_strategy("mem-123", "strat-1") + # Test getting specific strategy + 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() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + client, mock_client = _create_client() + + # Mock get_memory response (needed by update_memory_strategies for modify validation) + mock_client.get_memory.return_value = { + "memory": { + "id": "mem-123", + "strategies": [ + {"strategyId": "strat-456", "type": "SEMANTIC", "name": "Strategy 1"}, + ], + } + } - # Mock update_memory response (update_strategy uses update_memory internally) - mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} + # Mock update_memory response + 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"}, - ) + 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/"], + ) - 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 were passed - args, kwargs = mock_client.update_memory.call_args - assert kwargs["memoryId"] == "mem-123" - assert "memoryStrategies" in kwargs - assert "modifyMemoryStrategies" in kwargs["memoryStrategies"] + # 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"] - # 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"} + # 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/"] 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 - - error_response = {"Error": {"Code": "ValidationException", "Message": "Invalid parameter"}} - mock_client.create_memory.side_effect = ClientError(error_response, "CreateMemory") - - try: - client.create_memory(name="Test Memory") - raise AssertionError("Error was not raised as expected") - except ClientError as e: - assert "ValidationException" in str(e) + client, mock_client = _create_client() + error_response = {"Error": {"Code": "ValidationException", "Message": "Invalid parameter"}} + mock_client.create_memory.side_effect = ClientError(error_response, "CreateMemory") -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 - - # 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) - - assert result["id"] == "mem-123" - assert mock_client.get_memory.called - - # 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() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # 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")): - # 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 - - # 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" + client, mock_client = _create_client() + + # 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")): + # 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 + + # Verify parameters were passed (MemoryClient always includes memoryStrategies) + 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" + # Strategies may have been processed by _add_default_namespaces + assert "memoryStrategies" in kwargs + added_strategy = kwargs["memoryStrategies"][0] + assert "semanticMemoryStrategy" in added_strategy + assert added_strategy["semanticMemoryStrategy"]["name"] == "Strategy 1" def test_list_memories_with_pagination(): """Test list_memories with pagination.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + client, mock_client = _create_client() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_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)] - # 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)] + 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) - # Test with max_results requiring pagination - 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" - - # Verify two API calls were made - assert mock_client.list_memories.call_count == 2 + # Verify two API calls were made + 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" - - -def test_wait_for_status_timeout(): - """Test _wait_for_status with timeout.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client - - # 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"}]} - } + client, mock_client = _create_client() - # 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 response + 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")): + # Test minimal update (only memory_id) + result = client.update_memory(memory_id="mem-123") - # 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) + # Verify minimal parameters + 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 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"}], } + } - 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 delete_memory response + 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) + # Mock time to simulate timeout + 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_memory_active_timeout(): + """Test _wait_for_memory_active 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 get_memory to always return CREATING (never becomes 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 - ) + # Mock time to simulate timeout + 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_memory_active("mem-timeout", 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_memory_active_failure(): + """Test _wait_for_memory_active with FAILED status.""" + client, mock_client = _create_client() + # Mock get_memory to return FAILED status + 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_memory_active("mem-failed", 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_memory_active_with_strategy_check(): + """Test _wait_for_memory_active with transitional strategies.""" + client, mock_client = _create_client() - # Mock get_memory response without the newly added strategy - mock_client.get_memory.return_value = { + # Mock get_memory responses - first with transitional strategy, then all active + 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 +545,165 @@ def test_wait_for_status(): {"strategyId": "strat-2", "status": "ACTIVE"}, ], } + }, + ] + + with patch("time.time", return_value=0): + with patch("time.sleep"): + result = client._wait_for_memory_active("mem-123", max_wait=120, poll_interval=10) + + 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_add_strategy_strategy_not_found(): + """Test add_strategy when newly added strategy cannot be found for wait.""" + client, mock_client = _create_client() + + # Mock update_memory response (via update_memory_strategies) + mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} + + # Mock get_memory for _wait_for_memory_active - memory becomes ACTIVE with no strategies + mock_client.get_memory.return_value = { + "memory": { + "id": "mem-123", + "status": "ACTIVE", + "strategies": [], } + } + 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 _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 - ) + strategy = {"semanticMemoryStrategy": {"name": "Missing Strategy"}} + result = client.add_strategy("mem-123", strategy, wait_for_active=True) assert result["id"] == "mem-123" assert result["status"] == "ACTIVE" - assert mock_client.get_memory.called - - # Verify correct parameters - args, kwargs = mock_client.get_memory.call_args - assert kwargs["memoryId"] == "mem-123" 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() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + client, mock_client = _create_client() - # Mock ClientError - error_response = {"Error": {"Code": "AccessDeniedException", "Message": "Insufficient permissions"}} - mock_client.list_memories.side_effect = ClientError(error_response, "ListMemories") + 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() - - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + client, mock_client = _create_client() - # Mock ClientError - error_response = {"Error": {"Code": "ConflictException", "Message": "Memory is in use"}} - mock_client.delete_memory.side_effect = ClientError(error_response, "DeleteMemory") + 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_memory_active_client_error(): + """Test _wait_for_memory_active with ClientError.""" + client, mock_client = _create_client() -def test_wait_for_strategy_active_client_error(): - """Test _wait_for_strategy_active with ClientError.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # 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_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) + with patch("time.time", return_value=0): + with patch("time.sleep"): + try: + client._wait_for_memory_active("mem-123", max_wait=120, poll_interval=10) + raise AssertionError("ClientError was not raised") + except ClientError as e: + assert "InternalServerError" in str(e) -def test_wait_for_status_client_error(): - """Test _wait_for_status with ClientError.""" +def test_client_property_backward_compat(): + """Test that client.client provides backward-compatible access to the boto3 client.""" with patch("boto3.client"): client = MemoryControlPlaneClient() - # Mock the boto3 client - mock_client = MagicMock() - client.client = mock_client + # Setting client should update _memory_client.gmcp_client + mock_client = MagicMock() + client.client = mock_client + assert client.client is mock_client + assert client._memory_client.gmcp_client is mock_client - # Mock ClientError - error_response = {"Error": {"Code": "InternalServerError", "Message": "Internal server 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(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) +def test_remove_strategy(): + """Test remove_strategy functionality.""" + client, mock_client = _create_client() + + # Mock update_memory response (remove_strategy uses update_memory_strategies 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")): + result = client.remove_strategy(memory_id="mem-123", strategy_id="strat-456") + + assert result["id"] == "mem-123" + assert mock_client.update_memory.called + + # Verify delete strategy was passed correctly + 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"}] From fdfeb5b284fc9c10ce4055d07b04d3a2002ae496 Mon Sep 17 00:00:00 2001 From: "T.J Ariyawansa" Date: Tue, 10 Feb 2026 17:50:34 -0500 Subject: [PATCH 2/4] Deprecate MemoryControlPlaneClient with per-method warnings Keep original implementation (no pass-through to MemoryClient) to preserve exact return value contracts. Add deprecation warnings to __init__ and every public method with specific MemoryClient migration suggestions. Add client= constructor parameter for backward-compatible boto3 client injection. Closes #247 --- src/bedrock_agentcore/memory/controlplane.py | 441 ++++++++++++++---- .../memory/test_controlplane.py | 402 ++++++++-------- 2 files changed, 562 insertions(+), 281 deletions(-) diff --git a/src/bedrock_agentcore/memory/controlplane.py b/src/bedrock_agentcore/memory/controlplane.py index 333c23c7..ba539a29 100644 --- a/src/bedrock_agentcore/memory/controlplane.py +++ b/src/bedrock_agentcore/memory/controlplane.py @@ -5,19 +5,31 @@ """ 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 .client import MemoryClient 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. " + f"Use MemoryClient.{suggestion} instead. " + f"See: https://github.com/aws/bedrock-agentcore-sdk-python/issues/247", + DeprecationWarning, + stacklevel=3, + ) + + class MemoryControlPlaneClient: """Client for Bedrock AgentCore Memory control plane operations. @@ -26,12 +38,13 @@ class MemoryControlPlaneClient: 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 (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. " @@ -42,19 +55,18 @@ def __init__(self, region_name: str = "us-west-2", environment: str = "prod"): ) self.region_name = region_name self.environment = environment - self._memory_client = MemoryClient(region_name=region_name) - - logger.info("Initialized MemoryControlPlaneClient for %s in %s", environment, region_name) - @property - def client(self): - """The underlying boto3 control plane client.""" - return self._memory_client.gmcp_client + 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) - @client.setter - def client(self, value): - """Allow overriding the underlying boto3 client (used by tests).""" - self._memory_client.gmcp_client = value + logger.info("Initialized MemoryControlPlaneClient for %s in %s", environment, region_name) # ==================== MEMORY OPERATIONS ==================== @@ -71,6 +83,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) @@ -84,27 +99,47 @@ def create_memory( Returns: Created memory object """ - if wait_for_active: - return self._memory_client.create_memory_and_wait( - name=name, - strategies=strategies or [], - description=description, - event_expiry_days=event_expiry_days, - memory_execution_role_arn=memory_execution_role_arn, - max_wait=max_wait, - poll_interval=poll_interval, - ) - return self._memory_client.create_memory( - name=name, - strategies=strategies or [], - description=description, - event_expiry_days=event_expiry_days, - memory_execution_role_arn=memory_execution_role_arn, - ) + _deprecation("create_memory", "create_memory() or create_memory_and_wait()") + + params = { + "name": name, + "eventExpiryDuration": event_expiry_days, + "clientToken": str(uuid.uuid4()), + } + + if description: + params["description"] = description + + if memory_execution_role_arn: + params["memoryExecutionRoleArn"] = memory_execution_role_arn + + if strategies: + params["memoryStrategies"] = strategies + + try: + response = self.client.create_memory(**params) + memory = response["memory"] + memory_id = memory["id"] + + logger.info("Created memory: %s", memory_id) + + if wait_for_active: + return self._wait_for_memory_active(memory_id, max_wait, poll_interval) + + return memory + + except ClientError as e: + logger.error("Failed to create memory: %s", e) + raise 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 @@ -112,6 +147,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"] @@ -133,13 +170,43 @@ 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 """ - return self._memory_client.list_memories(max_results=max_results) + _deprecation("list_memories", "list_memories()") + + try: + memories = [] + next_token = None + + while len(memories) < max_results: + params = {"maxResults": min(100, max_results - len(memories))} + if next_token: + params["nextToken"] = next_token + + response = self.client.list_memories(**params) + batch = response.get("memories", []) + memories.extend(batch) + + next_token = response.get("nextToken") + if not next_token or len(memories) >= max_results: + break + + # Add strategy count to each memory summary + for memory in memories: + memory["strategyCount"] = 0 # List memories doesn't include strategies + + return memories[:max_results] + + except ClientError as e: + logger.error("Failed to list memories: %s", e) + raise def update_memory( self, @@ -156,6 +223,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 @@ -171,6 +242,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()), @@ -227,6 +303,9 @@ def delete_memory( ) -> 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 @@ -237,6 +316,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: @@ -253,30 +334,40 @@ 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, + target_status=MemoryStatus.ACTIVE.value, + max_wait=max_wait, + poll_interval=poll_interval, + check_strategies=True, ) - start_time = time.time() - while time.time() - start_time < max_wait: - memory = self.get_memory(memory_id) - strategies = memory.get("strategies", []) - all_ready = all( - s.get("status") in [MemoryStatus.ACTIVE.value, MemoryStatus.FAILED.value] - for s in strategies - ) - if all_ready: - break - time.sleep(poll_interval) except Exception as e: logger.warning("Error waiting for strategies to become ACTIVE: %s", e) - if wait_for_deletion: - return self._memory_client.delete_memory_and_wait( - memory_id=memory_id, - max_wait=max_wait, - poll_interval=poll_interval, - ) + # Now delete the memory + response = self.client.delete_memory(memoryId=memory_id, clientToken=str(uuid.uuid4())) + + logger.info("Initiated deletion of memory: %s", memory_id) - return self._memory_client.delete_memory(memory_id=memory_id) + if not wait_for_deletion: + return response + + # Wait for deletion to complete + start_time = time.time() + while time.time() - start_time < max_wait: + try: + self.client.get_memory(memoryId=memory_id) + time.sleep(poll_interval) + except ClientError as e: + if e.response["Error"]["Code"] == "ResourceNotFoundException": + logger.info("Memory %s successfully deleted", memory_id) + return response + raise + + raise TimeoutError(f"Memory {memory_id} was not deleted within {max_wait} seconds") except ClientError as e: logger.error("Failed to delete memory: %s", e) @@ -294,6 +385,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 @@ -302,23 +397,55 @@ def add_strategy( poll_interval: Seconds between status checks if wait_for_active is True Returns: - Updated memory object + Updated memory object with strategyId field """ - if wait_for_active: - return self._memory_client.update_memory_strategies_and_wait( - memory_id=memory_id, - add_strategies=[strategy], - max_wait=max_wait, - poll_interval=poll_interval, - ) - return self._memory_client.update_memory_strategies( + _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") + + logger.info("Adding strategy %s of type %s to memory %s", strategy_name, strategy_type, memory_id) + + # Use update_memory with add_strategies parameter but don't wait for memory + memory = self.update_memory( memory_id=memory_id, add_strategies=[strategy], + wait_for_active=False, # Don't wait for memory, we'll check strategy specifically ) + # If we need to wait for the strategy to become active + if wait_for_active: + # First, get the memory again to ensure we have the latest state + memory = self.get_memory(memory_id) + + # Find the newly added strategy by matching name + strategies = memory.get("strategies", []) + strategy_id = None + + for s in strategies: + # Match by name since that's unique within a memory + if s.get("name") == strategy_name: + strategy_id = s.get("strategyId") + logger.info("Found newly added strategy %s with ID %s", strategy_name, strategy_id) + break + + if strategy_id: + return self._wait_for_strategy_active(memory_id, strategy_id, max_wait, poll_interval) + else: + logger.warning("Could not identify newly added strategy %s to wait for activation", strategy_name) + + return memory + 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 @@ -326,8 +453,11 @@ def get_strategy(self, memory_id: str, strategy_id: str) -> Dict[str, Any]: Returns: Strategy details """ + _deprecation("get_strategy", "get_memory_strategies()") + try: - strategies = self._memory_client.get_memory_strategies(memory_id) + memory = self.get_memory(memory_id) + strategies = memory.get("strategies", []) for strategy in strategies: if strategy.get("strategyId") == strategy_id: @@ -352,6 +482,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 @@ -365,6 +499,12 @@ 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} if description is not None: @@ -376,18 +516,19 @@ def update_strategy( if configuration is not None: modify_config["configuration"] = configuration - if wait_for_active: - return self._memory_client.update_memory_strategies_and_wait( - memory_id=memory_id, - modify_strategies=[modify_config], - max_wait=max_wait, - poll_interval=poll_interval, - ) - return self._memory_client.update_memory_strategies( + # Use update_memory with modify_strategies parameter but don't wait for memory + memory = self.update_memory( memory_id=memory_id, modify_strategies=[modify_config], + wait_for_active=False, # Don't wait for memory, we'll check strategy specifically ) + # If we need to wait for the strategy to become active + if wait_for_active: + return self._wait_for_strategy_active(memory_id, strategy_id, max_wait, poll_interval) + + return memory + def remove_strategy( self, memory_id: str, @@ -398,6 +539,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 @@ -408,56 +553,162 @@ def remove_strategy( Returns: Updated memory object """ - if wait_for_active: - return self._memory_client.update_memory_strategies_and_wait( - memory_id=memory_id, - delete_strategy_ids=[strategy_id], - max_wait=max_wait, - poll_interval=poll_interval, - ) - return self._memory_client.update_memory_strategies( + _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( memory_id=memory_id, delete_strategy_ids=[strategy_id], + wait_for_active=wait_for_active, + max_wait=max_wait, + poll_interval=poll_interval, ) # ==================== HELPER METHODS ==================== def _wait_for_memory_active(self, memory_id: str, max_wait: int, poll_interval: int) -> Dict[str, Any]: - """Wait for memory and all strategies to reach ACTIVE state. + """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, + ) - Used by update_memory(wait_for_active=True). + def _wait_for_strategy_active( + self, memory_id: str, strategy_id: str, max_wait: int, poll_interval: int + ) -> Dict[str, Any]: + """Wait for specific memory strategy to become ACTIVE.""" + logger.info("Waiting for strategy %s to become ACTIVE (max wait: %d seconds)...", strategy_id, max_wait) + + start_time = time.time() + last_status = None + + while time.time() - start_time < max_wait: + try: + memory = self.get_memory(memory_id) + strategies = memory.get("strategies", []) + + for strategy in strategies: + if strategy.get("strategyId") == strategy_id: + status = strategy["status"] + + # Log status changes + if status != last_status: + logger.info("Strategy %s status: %s", strategy_id, status) + last_status = status + + if status == MemoryStatus.ACTIVE.value: + elapsed = time.time() - start_time + logger.info("Strategy %s is now ACTIVE (took %.1f seconds)", strategy_id, elapsed) + return memory + elif status == MemoryStatus.FAILED.value: + failure_reason = strategy.get("failureReason", "Unknown") + raise RuntimeError(f"Strategy {strategy_id} failed to activate: {failure_reason}") + + break + else: + logger.warning("Strategy %s not found in memory %s", strategy_id, memory_id) + + # Wait before checking again + time.sleep(poll_interval) + + except ClientError as e: + logger.error("Error checking strategy status: %s", e) + raise + + elapsed = time.time() - start_time + raise TimeoutError( + f"Strategy {strategy_id} did not become ACTIVE within {max_wait} seconds (last status: {last_status})" + ) + + def _wait_for_status( + 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. + + Args: + memory_id: The ID of the memory to check + target_status: The status to wait for (e.g., "ACTIVE") + max_wait: Maximum time to wait in seconds + poll_interval: Time between status checks in seconds + check_strategies: Whether to also check that all strategies are in the target status + + Returns: + The memory object once it reaches the target status + + Raises: + TimeoutError: If the memory doesn't reach the target status within max_wait + RuntimeError: If the memory or any strategy reaches a FAILED state """ - logger.info("Waiting for memory %s to become ACTIVE...", memory_id) + logger.info("Waiting for memory %s to reach status %s...", memory_id, target_status) start_time = time.time() last_memory_status = None + strategy_statuses = {} while time.time() - start_time < max_wait: try: memory = self.get_memory(memory_id) status = memory.get("status") + # Log status changes for memory if status != last_memory_status: logger.info("Memory %s status: %s", memory_id, status) last_memory_status = status - if status == MemoryStatus.ACTIVE.value: - strategies = memory.get("strategies", []) - all_strategies_active = all( - s.get("status") in [MemoryStatus.ACTIVE.value, MemoryStatus.FAILED.value] for s in strategies - ) - - if not all_strategies_active: - logger.info( - "Memory %s is ACTIVE but %d strategies are still processing", - memory_id, - len([s for s in strategies if s.get("status") != MemoryStatus.ACTIVE.value]), - ) - time.sleep(poll_interval) - continue + if status == target_status: + # Check if all strategies are also in the target status + if check_strategies and target_status == MemoryStatus.ACTIVE.value: + strategies = memory.get("strategies", []) + all_strategies_active = True + + for strategy in strategies: + strategy_id = strategy.get("strategyId") + strategy_status = strategy.get("status") + + # Log strategy status changes + if ( + strategy_id not in strategy_statuses + or strategy_statuses[strategy_id] != strategy_status + ): + logger.info("Strategy %s status: %s", strategy_id, strategy_status) + strategy_statuses[strategy_id] = strategy_status + + if strategy_status != target_status: + if strategy_status == MemoryStatus.FAILED.value: + failure_reason = strategy.get("failureReason", "Unknown") + raise RuntimeError(f"Strategy {strategy_id} failed: {failure_reason}") + + all_strategies_active = False + + if not all_strategies_active: + logger.info( + "Memory %s is %s but %d strategies are still processing", + memory_id, + target_status, + len([s for s in strategies if s.get("status") != target_status]), + ) + time.sleep(poll_interval) + continue elapsed = time.time() - start_time - logger.info("Memory %s and all strategies are now ACTIVE (took %.1f seconds)", memory_id, elapsed) + logger.info( + "Memory %s and all strategies are now %s (took %.1f seconds)", + memory_id, + target_status, + elapsed, + ) return memory elif status == MemoryStatus.FAILED.value: failure_reason = memory.get("failureReason", "Unknown") @@ -471,6 +722,6 @@ def _wait_for_memory_active(self, memory_id: str, max_wait: int, poll_interval: elapsed = time.time() - start_time raise TimeoutError( - f"Memory {memory_id} did not reach status ACTIVE within {max_wait} seconds " + f"Memory {memory_id} did not reach status {target_status} within {max_wait} seconds " f"(elapsed: {elapsed:.1f}s)" ) diff --git a/tests/bedrock_agentcore/memory/test_controlplane.py b/tests/bedrock_agentcore/memory/test_controlplane.py index 7565e28c..2a02e1d2 100644 --- a/tests/bedrock_agentcore/memory/test_controlplane.py +++ b/tests/bedrock_agentcore/memory/test_controlplane.py @@ -10,29 +10,67 @@ from bedrock_agentcore.memory.constants import MemoryStatus from bedrock_agentcore.memory.controlplane import MemoryControlPlaneClient -# Suppress the deprecation warning for all tests except the one that explicitly tests it +# 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 mocked boto3 clients.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() + """Helper to create a MemoryControlPlaneClient with a mocked boto3 client passed via constructor.""" mock_client = MagicMock() - client.client = mock_client + client = MemoryControlPlaneClient(client=mock_client) return client, mock_client def test_deprecation_warning(): - """Test that MemoryControlPlaneClient emits a deprecation warning.""" - with patch("boto3.client"): + """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") - MemoryControlPlaneClient() - assert len(w) == 1 - assert issubclass(w[0].category, DeprecationWarning) - assert "MemoryClient" in str(w[0].message) - assert "v1.4.0" in str(w[0].message) + 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(): @@ -45,7 +83,6 @@ def test_create_memory(): } with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - # Test basic memory creation (delegates to MemoryClient.create_memory) result = client.create_memory(name="Test Memory", description="Test description") assert result["id"] == "mem-123" @@ -76,7 +113,6 @@ def test_get_memory(): } } - # Test get memory with strategies result = client.get_memory("mem-123") assert result["id"] == "mem-123" @@ -91,18 +127,18 @@ def test_list_memories(): """Test list_memories functionality.""" client, mock_client = _create_client() - # 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} - # Test list memories (delegates to MemoryClient.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 # Verify API call args, kwargs = mock_client.list_memories.call_args @@ -113,19 +149,18 @@ def test_update_memory(): """Test update_memory functionality.""" client, mock_client = _create_client() - # 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")): - # Test memory update - result = client.update_memory(memory_id="mem-123", description="Updated description", event_expiry_days=120) + 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 - # Verify correct parameters args, kwargs = mock_client.update_memory.call_args assert kwargs["memoryId"] == "mem-123" assert kwargs["description"] == "Updated description" @@ -137,17 +172,14 @@ def test_delete_memory(): """Test delete_memory functionality.""" client, mock_client = _create_client() - # Mock response mock_client.delete_memory.return_value = {"status": "DELETING"} with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - # Test memory deletion (delegates to MemoryClient.delete_memory) result = client.delete_memory("mem-123") 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" @@ -157,22 +189,23 @@ def test_delete_memory_wait_for_strategies(): """Test delete_memory with wait_for_strategies=True.""" client, mock_client = _create_client() - # Mock get_memory responses: first with transitional strategies, then all active + # 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 = [ - # First call from get_memory (initial check) { "memory": { "id": "mem-123", + "status": "ACTIVE", "strategies": [ {"strategyId": "strat-1", "status": "CREATING"}, {"strategyId": "strat-2", "status": "ACTIVE"}, ], } }, - # Second call from get_memory (polling - all active now) { "memory": { "id": "mem-123", + "status": "ACTIVE", "strategies": [ {"strategyId": "strat-1", "status": "ACTIVE"}, {"strategyId": "strat-2", "status": "ACTIVE"}, @@ -181,7 +214,6 @@ def test_delete_memory_wait_for_strategies(): }, ] - # 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")): @@ -190,40 +222,28 @@ def test_delete_memory_wait_for_strategies(): result = client.delete_memory("mem-123", wait_for_strategies=True) assert result["status"] == "DELETING" - - # Verify get_memory was called to check strategy status - assert mock_client.get_memory.called - - # Verify delete_memory was called + assert mock_client.get_memory.call_count == 2 assert mock_client.delete_memory.called - args, kwargs = mock_client.delete_memory.call_args - assert kwargs["memoryId"] == "mem-123" def test_delete_memory_wait_for_deletion(): """Test delete_memory with wait_for_deletion=True.""" client, mock_client = _create_client() - # Mock delete_memory response mock_client.delete_memory.return_value = {"status": "DELETING"} - # Mock get_memory to raise ResourceNotFoundException (memory is gone) 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"): - result = client.delete_memory("mem-123", wait_for_deletion=True, max_wait=120, poll_interval=5) + 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" - - # Verify get_memory was called (to check if memory is gone) assert mock_client.get_memory.called @@ -231,10 +251,10 @@ def test_add_strategy(): """Test add_strategy functionality.""" client, mock_client = _create_client() - # Mock update_memory response (add_strategy uses update_memory_strategies 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"}} with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): @@ -243,7 +263,6 @@ def test_add_strategy(): assert result["id"] == "mem-123" assert mock_client.update_memory.called - # Verify strategy was passed in the update_memory call args, kwargs = mock_client.update_memory.call_args assert "memoryStrategies" in kwargs assert "addMemoryStrategies" in kwargs["memoryStrategies"] @@ -256,41 +275,52 @@ def test_add_strategy_wait_for_active(): """Test add_strategy with wait_for_active=True.""" client, mock_client = _create_client() - # Mock update_memory response (via update_memory_strategies) - mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} - - # Mock get_memory for _wait_for_memory_active (returns ACTIVE) - mock_client.get_memory.return_value = { - "memory": { - "id": "mem-123", - "status": "ACTIVE", - "strategies": [{"strategyId": "strat-new-123", "name": "Test Active Strategy", "status": "ACTIVE"}], - } + # First: update_memory response + mock_client.update_memory.return_value = { + "memory": {"id": "mem-123", "status": "CREATING"} } - # Test strategy addition with wait_for_active=True + # 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", + "status": "ACTIVE", + "strategies": [ + {"strategyId": "strat-new-123", "name": "Test Active Strategy", "status": "CREATING"}, + ], + } + }, + # _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"}, + ], + } + }, + ] + 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) + result = client.add_strategy( + "mem-123", strategy, wait_for_active=True, max_wait=120, poll_interval=5 + ) - # Result comes from _wait_for_memory_active which returns normalized memory 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"] - def test_get_strategy(): """Test get_strategy functionality.""" client, mock_client = _create_client() - # Mock get_memory response with strategies (used by get_memory_strategies) mock_client.get_memory.return_value = { "memory": { "id": "mem-123", @@ -301,7 +331,6 @@ def test_get_strategy(): } } - # Test getting specific strategy result = client.get_strategy("mem-123", "strat-1") assert result["strategyId"] == "strat-1" @@ -312,21 +341,11 @@ def test_update_strategy(): """Test update_strategy functionality.""" client, mock_client = _create_client() - # Mock get_memory response (needed by update_memory_strategies for modify validation) - mock_client.get_memory.return_value = { - "memory": { - "id": "mem-123", - "strategies": [ - {"strategyId": "strat-456", "type": "SEMANTIC", "name": "Strategy 1"}, - ], - } + mock_client.update_memory.return_value = { + "memory": {"id": "mem-123", "status": "CREATING"} } - # Mock update_memory response - 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", @@ -337,19 +356,39 @@ def test_update_strategy(): assert result["id"] == "mem-123" assert mock_client.update_memory.called - # 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"] - # 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/"] +def test_remove_strategy(): + """Test remove_strategy functionality.""" + client, mock_client = _create_client() + + 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")): + result = client.remove_strategy(memory_id="mem-123", strategy_id="strat-456") + + assert result["id"] == "mem-123" + assert mock_client.update_memory.called + + 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"}] + + def test_error_handling(): """Test error handling.""" client, mock_client = _create_client() @@ -368,13 +407,11 @@ def test_create_memory_with_strategies(): """Test create_memory with memory strategies.""" client, mock_client = _create_client() - # 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")): - # Test memory creation with strategies strategies = [{"semanticMemoryStrategy": {"name": "Strategy 1"}}] result = client.create_memory( name="Memory with Strategies", @@ -387,40 +424,37 @@ def test_create_memory_with_strategies(): assert result["id"] == "mem-456" assert mock_client.create_memory.called - # Verify parameters were passed (MemoryClient always includes memoryStrategies) 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" - # Strategies may have been processed by _add_default_namespaces - assert "memoryStrategies" in kwargs - added_strategy = kwargs["memoryStrategies"][0] - assert "semanticMemoryStrategy" in added_strategy - assert added_strategy["semanticMemoryStrategy"]["name"] == "Strategy 1" + assert kwargs["memoryStrategies"] == strategies def test_list_memories_with_pagination(): """Test list_memories with pagination.""" 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}, ] - # Test with max_results requiring pagination result = client.list_memories(max_results=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 @@ -428,17 +462,16 @@ def test_update_memory_minimal(): """Test update_memory with minimal parameters.""" client, mock_client = _create_client() - # Mock response - mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "ACTIVE"}} + 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" @@ -448,11 +481,12 @@ def test_get_strategy_not_found(): """Test get_strategy when strategy doesn't exist.""" 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"}], + "strategies": [ + {"strategyId": "strat-other", "name": "Other Strategy", "type": "SEMANTIC"} + ], } } @@ -467,64 +501,85 @@ def test_delete_memory_wait_for_deletion_timeout(): """Test delete_memory with wait_for_deletion timeout.""" client, mock_client = _create_client() - # Mock delete_memory response 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_client.get_memory.return_value = { + "memory": {"id": "mem-persistent", "status": "DELETING"} + } - # Mock time to simulate timeout 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")): + 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) + 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_memory_active_timeout(): - """Test _wait_for_memory_active with timeout.""" +def test_wait_for_status_timeout(): + """Test _wait_for_status with timeout.""" client, mock_client = _create_client() - # Mock get_memory to always return CREATING (never becomes ACTIVE) - mock_client.get_memory.return_value = {"memory": {"id": "mem-timeout", "status": "CREATING", "strategies": []}} + mock_client.get_memory.return_value = { + "memory": {"id": "mem-timeout", "status": "CREATING", "strategies": []} + } - # Mock time to simulate timeout 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_memory_active("mem-timeout", max_wait=300, poll_interval=10) + 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) -def test_wait_for_memory_active_failure(): - """Test _wait_for_memory_active with FAILED status.""" +def test_wait_for_status_failure(): + """Test _wait_for_status with FAILED status.""" client, mock_client = _create_client() - # Mock get_memory to return FAILED status mock_client.get_memory.return_value = { - "memory": {"id": "mem-failed", "status": "FAILED", "failureReason": "Configuration error", "strategies": []} + "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_memory_active("mem-failed", max_wait=300, poll_interval=10) + 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) -def test_wait_for_memory_active_with_strategy_check(): - """Test _wait_for_memory_active with transitional strategies.""" +def test_wait_for_status_with_strategy_check(): + """Test _wait_for_status with transitional strategies.""" client, mock_client = _create_client() - # Mock get_memory responses - first with transitional strategy, then all active mock_client.get_memory.side_effect = [ { "memory": { @@ -550,46 +605,25 @@ def test_wait_for_memory_active_with_strategy_check(): with patch("time.time", return_value=0): with patch("time.sleep"): - result = client._wait_for_memory_active("mem-123", max_wait=120, poll_interval=10) + result = client._wait_for_status( + "mem-123", + target_status=MemoryStatus.ACTIVE.value, + max_wait=120, + poll_interval=10, + ) 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_add_strategy_strategy_not_found(): - """Test add_strategy when newly added strategy cannot be found for wait.""" - client, mock_client = _create_client() - - # Mock update_memory response (via update_memory_strategies) - mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} - - # Mock get_memory for _wait_for_memory_active - memory becomes ACTIVE with no strategies - mock_client.get_memory.return_value = { - "memory": { - "id": "mem-123", - "status": "ACTIVE", - "strategies": [], - } - } - - with patch("uuid.uuid4", return_value=uuid.UUID("12345678-1234-5678-1234-567812345678")): - with patch("time.time", return_value=0): - with patch("time.sleep"): - strategy = {"semanticMemoryStrategy": {"name": "Missing Strategy"}} - result = client.add_strategy("mem-123", strategy, wait_for_active=True) - - assert result["id"] == "mem-123" - assert result["status"] == "ACTIVE" - - def test_get_memory_client_error(): """Test get_memory with ClientError.""" client, mock_client = _create_client() - error_response = {"Error": {"Code": "ResourceNotFoundException", "Message": "Memory not found"}} + error_response = { + "Error": {"Code": "ResourceNotFoundException", "Message": "Memory not found"} + } mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") try: @@ -603,7 +637,9 @@ def test_list_memories_client_error(): """Test list_memories with ClientError.""" client, mock_client = _create_client() - error_response = {"Error": {"Code": "AccessDeniedException", "Message": "Insufficient permissions"}} + error_response = { + "Error": {"Code": "AccessDeniedException", "Message": "Insufficient permissions"} + } mock_client.list_memories.side_effect = ClientError(error_response, "ListMemories") try: @@ -617,7 +653,9 @@ def test_update_memory_client_error(): """Test update_memory with ClientError.""" client, mock_client = _create_client() - error_response = {"Error": {"Code": "ValidationException", "Message": "Invalid memory parameters"}} + 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")): @@ -632,7 +670,9 @@ def test_delete_memory_client_error(): """Test delete_memory with ClientError.""" client, mock_client = _create_client() - error_response = {"Error": {"Code": "ConflictException", "Message": "Memory is in use"}} + 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")): @@ -647,7 +687,9 @@ def test_get_strategy_client_error(): """Test get_strategy with ClientError from get_memory.""" client, mock_client = _create_client() - error_response = {"Error": {"Code": "ThrottlingException", "Message": "Request throttled"}} + error_response = { + "Error": {"Code": "ThrottlingException", "Message": "Request throttled"} + } mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") try: @@ -657,53 +699,41 @@ def test_get_strategy_client_error(): assert "ThrottlingException" in str(e) -def test_wait_for_memory_active_client_error(): - """Test _wait_for_memory_active with ClientError.""" +def test_wait_for_status_client_error(): + """Test _wait_for_status with ClientError.""" client, mock_client = _create_client() - error_response = {"Error": {"Code": "InternalServerError", "Message": "Internal server error"}} + error_response = { + "Error": {"Code": "InternalServerError", "Message": "Internal server 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_memory_active("mem-123", max_wait=120, poll_interval=10) + 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) -def test_client_property_backward_compat(): - """Test that client.client provides backward-compatible access to the boto3 client.""" - with patch("boto3.client"): - client = MemoryControlPlaneClient() - - # Setting client should update _memory_client.gmcp_client +def test_constructor_client_param(): + """Test that passing client= in constructor sets self.client directly.""" mock_client = MagicMock() - client.client = mock_client - assert client.client is mock_client - assert client._memory_client.gmcp_client is mock_client + client = MemoryControlPlaneClient(client=mock_client) + assert client.client is mock_client -def test_remove_strategy(): - """Test remove_strategy functionality.""" - client, mock_client = _create_client() - - # Mock update_memory response (remove_strategy uses update_memory_strategies 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")): - 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_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() - # Verify delete strategy was passed correctly - 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"}] + assert client.client is mock_boto.return_value From 0c2c45c495d8e28593937aba1a19905b453c49ca Mon Sep 17 00:00:00 2001 From: "T.J Ariyawansa" Date: Tue, 10 Feb 2026 18:54:03 -0500 Subject: [PATCH 3/4] Remove issue URL from deprecation warning messages --- src/bedrock_agentcore/memory/controlplane.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/bedrock_agentcore/memory/controlplane.py b/src/bedrock_agentcore/memory/controlplane.py index ba539a29..f455f423 100644 --- a/src/bedrock_agentcore/memory/controlplane.py +++ b/src/bedrock_agentcore/memory/controlplane.py @@ -23,8 +23,7 @@ def _deprecation(method_name: str, suggestion: str) -> None: """Emit a deprecation warning for a MemoryControlPlaneClient method.""" warnings.warn( f"MemoryControlPlaneClient.{method_name}() is deprecated. " - f"Use MemoryClient.{suggestion} instead. " - f"See: https://github.com/aws/bedrock-agentcore-sdk-python/issues/247", + f"Use MemoryClient.{suggestion} instead.", DeprecationWarning, stacklevel=3, ) @@ -49,7 +48,7 @@ def __init__(self, region_name: str = "us-west-2", environment: str = "prod", cl 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. See: https://github.com/aws/bedrock-agentcore-sdk-python/issues/247", + "plus data plane features.", DeprecationWarning, stacklevel=2, ) From 0cb3d424a84ec0f5ac2bf6956528e496cf9e1ccf Mon Sep 17 00:00:00 2001 From: "T.J Ariyawansa" Date: Wed, 11 Feb 2026 11:37:49 -0500 Subject: [PATCH 4/4] Apply ruff formatting --- src/bedrock_agentcore/memory/controlplane.py | 3 +- .../memory/test_controlplane.py | 95 +++++-------------- 2 files changed, 25 insertions(+), 73 deletions(-) diff --git a/src/bedrock_agentcore/memory/controlplane.py b/src/bedrock_agentcore/memory/controlplane.py index f455f423..c6f7026b 100644 --- a/src/bedrock_agentcore/memory/controlplane.py +++ b/src/bedrock_agentcore/memory/controlplane.py @@ -22,8 +22,7 @@ def _deprecation(method_name: str, suggestion: str) -> None: """Emit a deprecation warning for a MemoryControlPlaneClient method.""" warnings.warn( - f"MemoryControlPlaneClient.{method_name}() is deprecated. " - f"Use MemoryClient.{suggestion} instead.", + f"MemoryControlPlaneClient.{method_name}() is deprecated. Use MemoryClient.{suggestion} instead.", DeprecationWarning, stacklevel=3, ) diff --git a/tests/bedrock_agentcore/memory/test_controlplane.py b/tests/bedrock_agentcore/memory/test_controlplane.py index 2a02e1d2..bcc6284e 100644 --- a/tests/bedrock_agentcore/memory/test_controlplane.py +++ b/tests/bedrock_agentcore/memory/test_controlplane.py @@ -36,16 +36,10 @@ 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.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.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} mock_client.delete_memory.return_value = {"status": "DELETING"} methods_and_suggestions = [ @@ -66,8 +60,7 @@ def test_method_deprecation_warnings(): 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) + 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) @@ -78,9 +71,7 @@ def test_create_memory(): client, mock_client = _create_client() # Mock successful response - mock_client.create_memory.return_value = { - "memory": {"id": "mem-123", "name": "Test Memory", "status": "CREATING"} - } + 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")): result = client.create_memory(name="Test Memory", description="Test description") @@ -154,9 +145,7 @@ def test_update_memory(): } 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 - ) + 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 @@ -238,9 +227,7 @@ def test_delete_memory_wait_for_deletion(): 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 - ) + result = client.delete_memory("mem-123", wait_for_deletion=True, max_wait=120, poll_interval=5) assert result["status"] == "DELETING" assert mock_client.delete_memory.called @@ -251,9 +238,7 @@ def test_add_strategy(): """Test add_strategy functionality.""" client, mock_client = _create_client() - mock_client.update_memory.return_value = { - "memory": {"id": "mem-123", "status": "CREATING"} - } + mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} strategy = {"semanticMemoryStrategy": {"name": "Test Strategy"}} @@ -276,9 +261,7 @@ def test_add_strategy_wait_for_active(): client, mock_client = _create_client() # First: update_memory response - mock_client.update_memory.return_value = { - "memory": {"id": "mem-123", "status": "CREATING"} - } + mock_client.update_memory.return_value = {"memory": {"id": "mem-123", "status": "CREATING"}} # Subsequent: get_memory calls for strategy lookup and wait mock_client.get_memory.side_effect = [ @@ -309,9 +292,7 @@ def test_add_strategy_wait_for_active(): 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 - ) + 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 @@ -341,9 +322,7 @@ def test_update_strategy(): """Test update_strategy functionality.""" client, mock_client = _create_client() - mock_client.update_memory.return_value = { - "memory": {"id": "mem-123", "status": "CREATING"} - } + 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")): result = client.update_strategy( @@ -371,9 +350,7 @@ def test_remove_strategy(): """Test remove_strategy functionality.""" client, mock_client = _create_client() - mock_client.update_memory.return_value = { - "memory": {"id": "mem-123", "status": "CREATING"} - } + 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")): result = client.remove_strategy(memory_id="mem-123", strategy_id="strat-456") @@ -436,12 +413,8 @@ def test_list_memories_with_pagination(): """Test list_memories with pagination.""" client, mock_client = _create_client() - 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"}, @@ -462,9 +435,7 @@ def test_update_memory_minimal(): """Test update_memory with minimal parameters.""" client, mock_client = _create_client() - mock_client.update_memory.return_value = { - "memory": {"id": "mem-123", "status": "ACTIVE"} - } + 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")): result = client.update_memory(memory_id="mem-123") @@ -484,9 +455,7 @@ def test_get_strategy_not_found(): mock_client.get_memory.return_value = { "memory": { "id": "mem-123", - "strategies": [ - {"strategyId": "strat-other", "name": "Other Strategy", "type": "SEMANTIC"} - ], + "strategies": [{"strategyId": "strat-other", "name": "Other Strategy", "type": "SEMANTIC"}], } } @@ -504,9 +473,7 @@ def test_delete_memory_wait_for_deletion_timeout(): 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_client.get_memory.return_value = {"memory": {"id": "mem-persistent", "status": "DELETING"}} with patch("time.time", side_effect=[0, 0, 0, 301, 301, 301]): with patch("time.sleep"): @@ -530,9 +497,7 @@ def test_wait_for_status_timeout(): """Test _wait_for_status with timeout.""" client, mock_client = _create_client() - mock_client.get_memory.return_value = { - "memory": {"id": "mem-timeout", "status": "CREATING", "strategies": []} - } + mock_client.get_memory.return_value = {"memory": {"id": "mem-timeout", "status": "CREATING", "strategies": []}} time_values = [0] + [i * 10 for i in range(1, 35)] + [301] with patch("time.time", side_effect=time_values): @@ -621,9 +586,7 @@ def test_get_memory_client_error(): """Test get_memory with ClientError.""" client, mock_client = _create_client() - error_response = { - "Error": {"Code": "ResourceNotFoundException", "Message": "Memory not found"} - } + error_response = {"Error": {"Code": "ResourceNotFoundException", "Message": "Memory not found"}} mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") try: @@ -637,9 +600,7 @@ def test_list_memories_client_error(): """Test list_memories with ClientError.""" client, mock_client = _create_client() - error_response = { - "Error": {"Code": "AccessDeniedException", "Message": "Insufficient permissions"} - } + error_response = {"Error": {"Code": "AccessDeniedException", "Message": "Insufficient permissions"}} mock_client.list_memories.side_effect = ClientError(error_response, "ListMemories") try: @@ -653,9 +614,7 @@ def test_update_memory_client_error(): """Test update_memory with ClientError.""" client, mock_client = _create_client() - error_response = { - "Error": {"Code": "ValidationException", "Message": "Invalid memory parameters"} - } + 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")): @@ -670,9 +629,7 @@ def test_delete_memory_client_error(): """Test delete_memory with ClientError.""" client, mock_client = _create_client() - error_response = { - "Error": {"Code": "ConflictException", "Message": "Memory is in use"} - } + 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")): @@ -687,9 +644,7 @@ def test_get_strategy_client_error(): """Test get_strategy with ClientError from get_memory.""" client, mock_client = _create_client() - error_response = { - "Error": {"Code": "ThrottlingException", "Message": "Request throttled"} - } + error_response = {"Error": {"Code": "ThrottlingException", "Message": "Request throttled"}} mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") try: @@ -703,9 +658,7 @@ def test_wait_for_status_client_error(): """Test _wait_for_status with ClientError.""" client, mock_client = _create_client() - error_response = { - "Error": {"Code": "InternalServerError", "Message": "Internal server error"} - } + error_response = {"Error": {"Code": "InternalServerError", "Message": "Internal server error"}} mock_client.get_memory.side_effect = ClientError(error_response, "GetMemory") with patch("time.time", return_value=0):