fix(identity): update endpoint for Create/UpdateWorkloadIdentity (#249)

Author: spatch1923

src/bedrock_agentcore/services/identity.py

@@ -78,9 +78,6 @@ def __init__(self, region: str):
         self.cp_client = boto3.client(
             "bedrock-agentcore-control", region_name=region, endpoint_url=get_control_plane_endpoint(region)
         )
-        self.identity_client = boto3.client(
-            "bedrock-agentcore-control", region_name=region, endpoint_url=get_data_plane_endpoint(region)
-        )
         self.dp_client = boto3.client(
             "bedrock-agentcore", region_name=region, endpoint_url=get_data_plane_endpoint(region)
         )
@@ -122,7 +119,7 @@ def create_workload_identity(
         self.logger.info("Creating workload identity...")
         if not name:
             name = f"workload-{uuid.uuid4().hex[:8]}"
-        return self.identity_client.create_workload_identity(
+        return self.cp_client.create_workload_identity(
             name=name, allowedResourceOauth2ReturnUrls=allowed_resource_oauth_2_return_urls or []
         )
 
@@ -131,7 +128,7 @@ def update_workload_identity(self, name: str, allowed_resource_oauth_2_return_ur
         self.logger.info(
             "Updating workload identity '%s' with callback urls: %s", name, allowed_resource_oauth_2_return_urls
         )
-        return self.identity_client.update_workload_identity(
+        return self.cp_client.update_workload_identity(
             name=name, allowedResourceOauth2ReturnUrls=allowed_resource_oauth_2_return_urls
         )
 

tests/bedrock_agentcore/services/test_identity.py

@@ -376,9 +376,8 @@ def test_get_workload_access_token_with_user_token(self):
 
         with patch("boto3.client") as mock_boto_client:
             mock_cp_client = Mock()
-            mock_identity_client = Mock()
             mock_dp_client = Mock()
-            mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_dp_client]
 
             identity_client = IdentityClient(region)
 
@@ -406,9 +405,8 @@ def test_get_workload_access_token_with_user_id(self):
 
         with patch("boto3.client") as mock_boto_client:
             mock_cp_client = Mock()
-            mock_identity_client = Mock()
             mock_dp_client = Mock()
-            mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_dp_client]
 
             identity_client = IdentityClient(region)
 
@@ -435,9 +433,8 @@ def test_get_workload_access_token_without_user_info(self):
 
         with patch("boto3.client") as mock_boto_client:
             mock_cp_client = Mock()
-            mock_identity_client = Mock()
             mock_dp_client = Mock()
-            mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_dp_client]
 
             identity_client = IdentityClient(region)
 
@@ -461,36 +458,35 @@ def test_create_workload_identity(self):
 
         with patch("boto3.client") as mock_boto_client:
             mock_cp_client = Mock()
-            mock_identity_client = Mock()
             mock_dp_client = Mock()
-            mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_dp_client]
 
             identity_client = IdentityClient(region)
 
             # Test with provided name
             custom_name = "my-custom-workload"
             expected_response = {"name": custom_name, "workloadIdentityId": "workload-123"}
-            mock_identity_client.create_workload_identity.return_value = expected_response
+            mock_cp_client.create_workload_identity.return_value = expected_response
 
             result = identity_client.create_workload_identity(name=custom_name)
 
             assert result == expected_response
-            mock_identity_client.create_workload_identity.assert_called_with(
+            mock_cp_client.create_workload_identity.assert_called_with(
                 name=custom_name, allowedResourceOauth2ReturnUrls=[]
             )
 
             # Test without provided name (auto-generated)
-            mock_identity_client.reset_mock()
+            mock_cp_client.reset_mock()
             expected_response_auto = {"name": "workload-abcd1234", "workloadIdentityId": "workload-456"}
-            mock_identity_client.create_workload_identity.return_value = expected_response_auto
+            mock_cp_client.create_workload_identity.return_value = expected_response_auto
 
             with patch("uuid.uuid4") as mock_uuid:
                 mock_uuid.return_value.hex = "abcd1234efgh5678"
 
                 result = identity_client.create_workload_identity()
 
                 assert result == expected_response_auto
-                mock_identity_client.create_workload_identity.assert_called_with(
+                mock_cp_client.create_workload_identity.assert_called_with(
                     name="workload-abcd1234", allowedResourceOauth2ReturnUrls=[]
                 )
 
@@ -499,22 +495,21 @@ def test_update_workload_identity(self):
 
         with patch("boto3.client") as mock_boto_client:
             mock_cp_client = Mock()
-            mock_identity_client = Mock()
             mock_dp_client = Mock()
-            mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_dp_client]
 
             identity_client = IdentityClient(region)
 
             workload_name = "test-workload"
             allowed_urls = ["https://unit-test.com/callback", "https://test.com/oauth"]
             expected_response = {"name": workload_name, "allowedResourceOauth2ReturnUrls": allowed_urls}
 
-            mock_identity_client.update_workload_identity.return_value = expected_response
+            mock_cp_client.update_workload_identity.return_value = expected_response
 
             result = identity_client.update_workload_identity(workload_name, allowed_urls)
 
             assert result == expected_response
-            mock_identity_client.update_workload_identity.assert_called_once_with(
+            mock_cp_client.update_workload_identity.assert_called_once_with(
                 name=workload_name, allowedResourceOauth2ReturnUrls=allowed_urls
             )
 
@@ -523,9 +518,8 @@ def test_get_workload_identity(self):
 
         with patch("boto3.client") as mock_boto_client:
             mock_cp_client = Mock()
-            mock_identity_client = Mock()
             mock_dp_client = Mock()
-            mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_dp_client]
 
             identity_client = IdentityClient(region)
 
@@ -545,9 +539,8 @@ def test_complete_resource_token_auth_with_user_id(self):
 
         with patch("boto3.client") as mock_boto_client:
             mock_cp_client = Mock()
-            mock_identity_client = Mock()
             mock_dp_client = Mock()
-            mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_dp_client]
 
             identity_client = IdentityClient(region)
 
@@ -570,9 +563,8 @@ def test_complete_resource_token_auth_with_user_token(self):
 
         with patch("boto3.client") as mock_boto_client:
             mock_cp_client = Mock()
-            mock_identity_client = Mock()
             mock_dp_client = Mock()
-            mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_dp_client]
 
             identity_client = IdentityClient(region)
 
@@ -595,9 +587,8 @@ def test_complete_resource_token_auth_with_invalid_identifier(self):
 
         with patch("boto3.client") as mock_boto_client:
             mock_cp_client = Mock()
-            mock_identity_client = Mock()
             mock_dp_client = Mock()
-            mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_dp_client]
 
             identity_client = IdentityClient(region)