Merge pull request #10 from aws/updateRegionnMemory

Update region in memory client

Author: sundargthb

src/bedrock_agentcore/memory/client.py

@@ -10,7 +10,6 @@
 
 import copy
 import logging
-import os
 import time
 import uuid
 import warnings
@@ -39,39 +38,19 @@
 class MemoryClient:
     """High-level Bedrock AgentCore Memory client with essential operations."""
 
-    def __init__(self, region_name: str = "us-west-2", environment: str = "prod"):
+    def __init__(self, region_name: Optional[str] = None):
         """Initialize the Memory client."""
-        self.region_name = region_name
-        if os.getenv("AWS_REGION"):
-            env_region = os.getenv("AWS_REGION")
-            if env_region != region_name:
-                warnings.warn(
-                    f"AWS_REGION environment variable ({env_region}) differs from provided "
-                    f"region_name ({region_name}). Using provided region_name.",
-                    stacklevel=2,
-                )
-        self.environment = environment
-
-        self.control_plane_endpoint = os.getenv(
-            "AGENTCORE_CONTROL_ENDPOINT", "https://bedrock-agentcore-control.us-west-2.amazonaws.com"
-        )
-        self.data_plane_endpoint = os.getenv(
-            "AGENTCORE_DATA_ENDPOINT", "https://bedrock-agentcore.us-west-2.amazonaws.com"
-        )
+        self.region_name = boto3.Session().region_name or region_name or "us-west-2"
 
-        control_service = os.getenv("AGENTCORE_CONTROL_SERVICE", "bedrock-agentcore-control")
-        data_service = os.getenv("AGENTCORE_DATA_SERVICE", "bedrock-agentcore")
+        self.gmcp_client = boto3.client("bedrock-agentcore-control", region_name=self.region_name)
+        self.gmdp_client = boto3.client("bedrock-agentcore", region_name=self.region_name)
 
-        self.gmcp_client = boto3.client(
-            control_service, region_name=self.region_name, endpoint_url=self.control_plane_endpoint
+        logger.info(
+            "Initialized MemoryClient for control plane: %s, data plane: %s",
+            self.gmcp_client.meta.region_name,
+            self.gmdp_client.meta.region_name,
         )
 
-        self.gmdp_client = boto3.client(
-            data_service, region_name=self.region_name, endpoint_url=self.data_plane_endpoint
-        )
-
-        logger.info("Initialized MemoryClient for %s in %s", environment, region_name)
-
     def create_memory(
         self,
         name: str,

tests/bedrock_agentcore/memory/test_client.py

@@ -14,36 +14,53 @@
 def test_client_initialization():
     """Test client initialization."""
     with patch("boto3.client") as mock_boto_client:
-        # Test gamma environment
-        client = MemoryClient(region_name="us-west-2", environment="gamma")
+        # Setup the mock to return a consistent region_name
+        mock_client_instance = MagicMock()
+        mock_client_instance.meta.region_name = "us-west-2"
+        mock_boto_client.return_value = mock_client_instance
 
+        client = MemoryClient(region_name="us-west-2")
+
+        # Check that the region was set correctly and boto3.client was called twice
         assert client.region_name == "us-west-2"
         assert mock_boto_client.call_count == 2
 
 
 def test_client_initialization_region_mismatch():
     """Test client initialization with region mismatch warning."""
-    with patch("boto3.client"):
-        with patch("os.getenv") as mock_getenv:
-            # Mock AWS_REGION environment variable
-            mock_getenv.return_value = "us-east-1"
-
-            with warnings.catch_warnings(record=True) as w:
-                warnings.simplefilter("always")
-
-                # Initialize client with different region than AWS_REGION
-                client = MemoryClient(region_name="us-west-2", environment="prod")
 
-                # Verify warning was issued
-                assert len(w) >= 1
-                warning_messages = [str(warning.message) for warning in w]
-                assert any(
-                    "AWS_REGION environment variable (us-east-1) differs from provided region_name (us-west-2)" in msg
-                    for msg in warning_messages
-                )
-
-                # Verify client uses provided region_name, not environment variable
-                assert client.region_name == "us-west-2"
+    with patch("boto3.client") as mock_boto_client:
+        # First test - environment variable takes precedence
+        with patch("boto3.Session") as mock_session:
+            # Mock the session instance to simulate AWS_REGION=us-east-1
+            mock_session_instance = MagicMock()
+            mock_session_instance.region_name = "us-east-1"
+            mock_session.return_value = mock_session_instance
+
+            # Mock the boto client
+            mock_client_instance = MagicMock()
+            mock_client_instance.meta.region_name = "us-east-1"
+            mock_boto_client.return_value = mock_client_instance
+
+            # When region_name is provided, environment variable should still take precedence
+            client1 = MemoryClient(region_name="us-west-2")
+            assert client1.region_name == "us-east-1"  # Environment wins over explicit param
+
+        # Second test - no environment variable, explicit param is used
+        with patch("boto3.Session") as mock_session:
+            # Mock the session instance to simulate no AWS_REGION set
+            mock_session_instance = MagicMock()
+            mock_session_instance.region_name = None
+            mock_session.return_value = mock_session_instance
+
+            # Mock the boto client
+            mock_client_instance = MagicMock()
+            mock_client_instance.meta.region_name = "us-west-2"
+            mock_boto_client.return_value = mock_client_instance
+
+            # When AWS_REGION is not set, explicitly provided region should be used
+            client2 = MemoryClient(region_name="us-west-2")
+            assert client2.region_name == "us-west-2"
 
 
 def test_namespace_defaults():

tests_integ/memory/test_devex.py

@@ -690,7 +690,7 @@ def main():
 
     logger.info("Using region: %s, environment: %s", region, environment)
 
-    client = MemoryClient(region_name=region, environment=environment)
+    client = MemoryClient(region_name=region)
 
     logger.info("\nCreating test memory with strategies...")
     memory = client.create_memory(

tests_integ/memory/test_memory_client.py

@@ -348,7 +348,7 @@ def main():
 
     logger.info("Using region: %s, environment: %s", region, environment)
 
-    client = MemoryClient(region_name=region, environment=environment)
+    client = MemoryClient(region_name=region)
 
     # Test Issue #2 first (strategy polling)
     test_strategy_polling_fix(client)