remove environment variables from client

Author: “Sundar”

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,15 @@
 class MemoryClient:
     """High-level Bedrock AgentCore Memory client with essential operations."""
 
-    def __init__(self, region_name: Optional[str] = None, environment: str = "prod"):
+    def __init__(self, region_name: Optional[str] = None):
         """Initialize the Memory client."""
-        env_region = os.getenv("AWS_REGION")
-        if region_name is not None and env_region is not None and 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.region_name = region_name
-        else:
-            self.region_name = region_name or env_region or boto3.Session().region_name or "us-west-2"
-        self.environment = environment
-
-        self.control_plane_endpoint = os.getenv(
-            "AGENTCORE_CONTROL_ENDPOINT", f"https://bedrock-agentcore-control.{self.region_name}.amazonaws.com"
-        )
-        self.data_plane_endpoint = os.getenv(
-            "AGENTCORE_DATA_ENDPOINT", f"https://bedrock-agentcore.{self.region_name}.amazonaws.com"
-        )
+        self.gmcp_client = boto3.client("bedrock-agentcore-control", region_name=region_name)
+        self.gmdp_client = boto3.client("bedrock-agentcore", region_name=region_name)
 
-        control_service = os.getenv("AGENTCORE_CONTROL_SERVICE", "bedrock-agentcore-control")
-        data_service = os.getenv("AGENTCORE_DATA_SERVICE", "bedrock-agentcore")
-
-        self.gmcp_client = boto3.client(
-            control_service, region_name=self.region_name, endpoint_url=self.control_plane_endpoint
-        )
-
-        self.gmdp_client = boto3.client(
-            data_service, region_name=self.region_name, endpoint_url=self.data_plane_endpoint
-        )
+        # Store the region that was actually used
+        self.region_name = self.gmcp_client.meta.region_name
 
-        logger.info("Initialized MemoryClient for %s in %s", environment, region_name)
+        logger.info("Initialized MemoryClient for %s", region_name)
 
     def create_memory(
         self,

tests/bedrock_agentcore/memory/test_client.py

@@ -14,36 +14,40 @@
 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 mock a default region when none is specified
+        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 not provided, it should use the boto3 default
+        client1 = MemoryClient()
+        assert client1.region_name == "us-east-1"
+
+        # Reset the mock to test with a specified region
+        mock_boto_client.reset_mock()
+        mock_client_instance = MagicMock()
+        mock_client_instance.meta.region_name = "us-west-2"
+        mock_boto_client.return_value = mock_client_instance
+
+        # When region_name is provided, it should use that value
+        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)