From 72a55fa8d528f505523a055eeb2180f6bc8ed66d Mon Sep 17 00:00:00 2001 From: Ken Guo Date: Fri, 1 Aug 2025 12:26:00 -0700 Subject: [PATCH] use json to manage local workload identity and user id --- .gitignore | 1 + src/bedrock_agentcore/identity/auth.py | 9 +- tests/bedrock_agentcore/identity/test_auth.py | 94 ++++++++++--------- tests_integ/identity/test_auth_flows.py | 7 +- 4 files changed, 60 insertions(+), 51 deletions(-) diff --git a/.gitignore b/.gitignore index 27beb74..1330d9c 100644 --- a/.gitignore +++ b/.gitignore @@ -167,6 +167,7 @@ cython_debug/ # macOS .DS_Store .agentcore.yaml +.agentcore.json .AppleDouble .LSOverride diff --git a/src/bedrock_agentcore/identity/auth.py b/src/bedrock_agentcore/identity/auth.py index a2024e5..ca7d9c0 100644 --- a/src/bedrock_agentcore/identity/auth.py +++ b/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") diff --git a/tests/bedrock_agentcore/identity/test_auth.py b/tests/bedrock_agentcore/identity/test_auth.py index 89b915c..34c0777 100644 --- a/tests/bedrock_agentcore/identity/test_auth.py +++ b/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: diff --git a/tests_integ/identity/test_auth_flows.py b/tests_integ/identity/test_auth_flows.py index 1874329..70925c2 100644 --- a/tests_integ/identity/test_auth_flows.py +++ b/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=""))