use dp endpoint to create workload identity

Author: guoken1

src/bedrock_agentcore/services/identity.py

@@ -65,6 +65,9 @@ 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)
         )
@@ -104,7 +107,7 @@ def create_workload_identity(self, name: Optional[str] = None) -> Dict:
         self.logger.info("Creating workload identity...")
         if not name:
             name = f"workload-{uuid.uuid4().hex[:8]}"
-        return self.cp_client.create_workload_identity(name=name)
+        return self.identity_client.create_workload_identity(name=name)
 
     async def get_token(
         self,

tests/bedrock_agentcore/services/test_identity.py

@@ -328,8 +328,9 @@ 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_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client]
 
             identity_client = IdentityClient(region)
 
@@ -357,8 +358,9 @@ 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_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client]
 
             identity_client = IdentityClient(region)
 
@@ -385,8 +387,9 @@ 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_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client]
 
             identity_client = IdentityClient(region)
 
@@ -410,33 +413,34 @@ 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_dp_client]
+            mock_boto_client.side_effect = [mock_cp_client, mock_identity_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_cp_client.create_workload_identity.return_value = expected_response
+            mock_identity_client.create_workload_identity.return_value = expected_response
 
             result = identity_client.create_workload_identity(name=custom_name)
 
             assert result == expected_response
-            mock_cp_client.create_workload_identity.assert_called_with(name=custom_name)
+            mock_identity_client.create_workload_identity.assert_called_with(name=custom_name)
 
             # Test without provided name (auto-generated)
-            mock_cp_client.reset_mock()
+            mock_identity_client.reset_mock()
             expected_response_auto = {"name": "workload-abcd1234", "workloadIdentityId": "workload-456"}
-            mock_cp_client.create_workload_identity.return_value = expected_response_auto
+            mock_identity_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_cp_client.create_workload_identity.assert_called_with(name="workload-abcd1234")
+                mock_identity_client.create_workload_identity.assert_called_with(name="workload-abcd1234")
 
 
 class TestDefaultApiTokenPoller: