use json to manage local workload identity and user id (#37)

Author: guoken1

.gitignore

@@ -167,6 +167,7 @@ cython_debug/
 # macOS
 .DS_Store
 .agentcore.yaml
+.agentcore.json
 .AppleDouble
 .LSOverride
 

src/bedrock_agentcore/identity/auth.py

@@ -157,18 +157,17 @@ async def _get_workload_access_token(client: IdentityClient) -> str:
 
 
 async def _set_up_local_auth(client: IdentityClient) -> str:
+    import json
     import uuid
     from pathlib import Path
 
-    import yaml
-
-    config_path = Path(".agentcore.yaml")
+    config_path = Path(".agentcore.json")
     workload_identity_name = None
     config = {}
     if config_path.exists():
         try:
             with open(config_path, "r", encoding="utf-8") as file:
-                config = yaml.safe_load(file) or {}
+                config = json.load(file) or {}
         except Exception:
             print("Could not find existing workload identity and user id")
 
@@ -189,7 +188,7 @@ async def _set_up_local_auth(client: IdentityClient) -> str:
     try:
         config = {"workload_identity_name": workload_identity_name, "user_id": user_id}
         with open(config_path, "w", encoding="utf-8") as file:
-            yaml.dump(config, file, default_flow_style=False, indent=2)
+            json.dump(config, file, indent=2)
     except Exception:
         print("Warning: could not write the created workload identity to file")
 

tests/bedrock_agentcore/identity/test_auth.py

@@ -1,7 +1,8 @@
 """Tests for Bedrock AgentCore authentication decorators and functions."""
 
+import json
 import os
-from unittest.mock import AsyncMock, Mock, mock_open, patch
+from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 
@@ -332,60 +333,67 @@ class TestSetUpLocalAuth:
     """Test _set_up_local_auth function."""
 
     @pytest.mark.asyncio
-    async def test_existing_config(self):
+    async def test_existing_config(self, tmp_path):
         """Test when config file exists with both workload_identity_name and user_id."""
         config_content = {"workload_identity_name": "existing-workload-123", "user_id": "existing-user-456"}
         mock_client = Mock()
         mock_client.get_workload_access_token = Mock(return_value={"workloadAccessToken": "test-access-token-456"})
 
-        with patch("pathlib.Path") as mock_path_class:
-            mock_path = Mock()
-            mock_path.exists.return_value = True
-            mock_path.absolute.return_value = "/test/.agentcore.yaml"
-            mock_path_class.return_value = mock_path
-
-            with patch("builtins.open", mock_open()):
-                with patch("yaml.safe_load", return_value=config_content):
-                    result = await _set_up_local_auth(mock_client)
-
-                    # Should use existing workload identity and user_id
-                    assert result == "test-access-token-456"
-                    mock_client.create_workload_identity.assert_not_called()
-                    mock_client.get_workload_access_token.assert_called_once_with(
-                        "existing-workload-123", user_id="existing-user-456"
-                    )
+        # Create the config file in the temp directory
+        config_file = tmp_path / ".agentcore.json"
+        config_file.write_text(json.dumps(config_content))
+
+        # Change to the temp directory for the test
+        import os
+
+        original_dir = os.getcwd()
+        try:
+            os.chdir(tmp_path)
+            result = await _set_up_local_auth(mock_client)
+
+            # Should use existing workload identity and user_id
+            assert result == "test-access-token-456"
+            mock_client.create_workload_identity.assert_not_called()
+            mock_client.get_workload_access_token.assert_called_once_with(
+                "existing-workload-123", user_id="existing-user-456"
+            )
+        finally:
+            os.chdir(original_dir)
 
     @pytest.mark.asyncio
-    async def test_no_config(self):
+    async def test_no_config(self, tmp_path):
         """Test when config file doesn't exist."""
         mock_client = Mock()
         mock_client.create_workload_identity = Mock(return_value={"name": "test-workload-123"})
         mock_client.get_workload_access_token = Mock(return_value={"workloadAccessToken": "test-access-token-456"})
 
-        with patch("pathlib.Path") as mock_path_class:
-            mock_path = Mock()
-            mock_path.exists.return_value = False
-            mock_path_class.return_value = mock_path
-
-            with patch("builtins.open", mock_open()):
-                with patch("yaml.dump") as mock_yaml_dump:
-                    with patch("uuid.uuid4") as mock_uuid:
-                        mock_uuid.return_value.hex = "abcd1234efgh5678"
-
-                        result = await _set_up_local_auth(mock_client)
-
-                        # Should create new workload identity and user_id
-                        assert result == "test-access-token-456"
-                        mock_client.create_workload_identity.assert_called_once()
-                        mock_client.get_workload_access_token.assert_called_once_with(
-                            "test-workload-123", user_id="abcd1234"
-                        )
-
-                        # Should create and save new config
-                        mock_yaml_dump.assert_called_once()
-                        saved_config = mock_yaml_dump.call_args[0][0]
-                        assert saved_config["workload_identity_name"] == "test-workload-123"
-                        assert saved_config["user_id"] == "abcd1234"
+        # Change to the temp directory for the test
+        import os
+
+        original_dir = os.getcwd()
+        try:
+            os.chdir(tmp_path)
+
+            with patch("uuid.uuid4") as mock_uuid:
+                mock_uuid.return_value.hex = "abcd1234efgh5678"
+
+                result = await _set_up_local_auth(mock_client)
+
+                # Should create new workload identity and user_id
+                assert result == "test-access-token-456"
+                mock_client.create_workload_identity.assert_called_once()
+                mock_client.get_workload_access_token.assert_called_once_with("test-workload-123", user_id="abcd1234")
+
+                # Verify that the config file was created
+                config_file = tmp_path / ".agentcore.json"
+                assert config_file.exists()
+
+                # Verify the config file content
+                saved_config = json.loads(config_file.read_text())
+                assert saved_config["workload_identity_name"] == "test-workload-123"
+                assert saved_config["user_id"] == "abcd1234"
+        finally:
+            os.chdir(original_dir)
 
 
 class TestGetRegion:

tests_integ/identity/test_auth_flows.py

@@ -30,6 +30,7 @@ async def need_api_key(*, api_key: str):
     print(f"received api key for async func: {api_key}")
 
 
-asyncio.run(need_api_key(api_key=""))
-asyncio.run(need_token_2LO_async(access_token=""))
-asyncio.run(need_token_3LO_async(access_token=""))
+if __name__ == "__main__":
+    asyncio.run(need_api_key(api_key=""))
+    asyncio.run(need_token_2LO_async(access_token=""))
+    asyncio.run(need_token_3LO_async(access_token=""))