From f27320b3eec9060061f6dc7aa5e13713678552ff Mon Sep 17 00:00:00 2001 From: Guillaume Charest <1690085+gcharest@users.noreply.github.com> Date: Fri, 19 Apr 2024 17:27:54 +0000 Subject: [PATCH] feat: add create, delete, get user functions --- app/integrations/aws/identity_store.py | 62 ++++++++ .../integrations/aws/test_identity_store.py | 138 ++++++++++++++++++ 2 files changed, 200 insertions(+) diff --git a/app/integrations/aws/identity_store.py b/app/integrations/aws/identity_store.py index c6f575bb..f35c06d4 100644 --- a/app/integrations/aws/identity_store.py +++ b/app/integrations/aws/identity_store.py @@ -23,6 +23,68 @@ def resolve_identity_store_id(kwargs): return kwargs +@handle_aws_api_errors +def create_user(email, first_name, family_name, **kwargs): + """Creates a new user in the AWS Identity Center (identitystore) + + Args: + email (str): The email address of the user. + first_name (str): The first name of the user. + family_name (str): The family name of the user. + **kwargs: Additional keyword arguments for the API call. + + Returns: + str: The user ID of the created user. + """ + kwargs = resolve_identity_store_id(kwargs) + kwargs.update( + { + "UserName": email, + "Emails": [{"Value": email, "Type": "WORK", "Primary": True}], + "Name": {"GivenName": first_name, "FamilyName": family_name}, + "DisplayName": f"{first_name} {family_name}", + } + ) + return execute_aws_api_call("identitystore", "create_user", **kwargs)["UserId"] + + +@handle_aws_api_errors +def delete_user(user_id, **kwargs): + """Deletes a user from the AWS Identity Center (identitystore) + + Args: + user_id (str): The user ID of the user. + **kwargs: Additional keyword arguments for the API call. + """ + kwargs = resolve_identity_store_id(kwargs) + kwargs.update({"UserId": user_id}) + result = execute_aws_api_call("identitystore", "delete_user", **kwargs) + return True if result == {} else False + + +@handle_aws_api_errors +def get_user_id(user_name, **kwargs): + """Retrieves the user ID of the current user + + Args: + user_name (str): The user name of the user. Default is the primary email address. + **kwargs: Additional keyword arguments for the API call. + """ + kwargs = resolve_identity_store_id(kwargs) + kwargs.update( + { + "AlternateIdentifier": { + "UniqueAttribute": { + "AttributePath": "userName", + "AttributeValue": user_name, + }, + } + } + ) + result = execute_aws_api_call("identitystore", "get_user_id", **kwargs) + return result["UserId"] if result else False + + @handle_aws_api_errors def list_users(**kwargs): """Retrieves all users from the AWS Identity Center (identitystore)""" diff --git a/app/tests/integrations/aws/test_identity_store.py b/app/tests/integrations/aws/test_identity_store.py index 530ed277..cda4df3e 100644 --- a/app/tests/integrations/aws/test_identity_store.py +++ b/app/tests/integrations/aws/test_identity_store.py @@ -23,6 +23,144 @@ def test_resolve_identity_store_id_no_env(): identity_store.resolve_identity_store_id({}) +@patch("integrations.aws.identity_store.execute_aws_api_call") +@patch("integrations.aws.identity_store.resolve_identity_store_id") +def test_create_user(mock_resolve_identity_store_id, mock_execute_aws_api_call): + mock_resolve_identity_store_id.return_value = { + "IdentityStoreId": "test_instance_id" + } + mock_execute_aws_api_call.return_value = {"UserId": "test_user_id"} + email = "test@example.com" + first_name = "Test" + family_name = "User" + + # Act + result = identity_store.create_user(email, first_name, family_name) + + # Assert + mock_execute_aws_api_call.assert_called_once_with( + "identitystore", + "create_user", + IdentityStoreId="test_instance_id", + UserName=email, + Emails=[{"Value": email, "Type": "WORK", "Primary": True}], + Name={"GivenName": first_name, "FamilyName": family_name}, + DisplayName=f"{first_name} {family_name}", + ) + assert result == "test_user_id" + + +@patch("integrations.aws.identity_store.execute_aws_api_call") +@patch("integrations.aws.identity_store.resolve_identity_store_id") +def test_get_user_id(mock_resolve_identity_store_id, mock_execute_aws_api_call): + mock_resolve_identity_store_id.return_value = { + "IdentityStoreId": "test_instance_id" + } + mock_execute_aws_api_call.return_value = {"UserId": "test_user_id"} + email = "test@example.com" + user_name = email + request = { + "AlternateIdentifier": { + "UniqueAttribute": { + "AttributePath": "userName", + "AttributeValue": user_name, + }, + }, + } + + # Act + result = identity_store.get_user_id(user_name) + + # Assert + mock_execute_aws_api_call.assert_called_once_with( + "identitystore", + "get_user_id", + IdentityStoreId="test_instance_id", + **request, + ) + assert result == "test_user_id" + + +@patch("integrations.aws.identity_store.execute_aws_api_call") +@patch("integrations.aws.identity_store.resolve_identity_store_id") +def test_get_user_id_user_not_found( + mock_resolve_identity_store_id, mock_execute_aws_api_call +): + # Arrange + mock_resolve_identity_store_id.return_value = { + "IdentityStoreId": "test_instance_id" + } + mock_execute_aws_api_call.return_value = False + user_name = "nonexistent_user" + + # Act + result = identity_store.get_user_id(user_name) + + # Assert + mock_execute_aws_api_call.assert_called_once_with( + "identitystore", + "get_user_id", + IdentityStoreId="test_instance_id", + AlternateIdentifier={ + "UniqueAttribute": { + "AttributePath": "userName", + "AttributeValue": user_name, + }, + }, + ) + assert result is False + + +@patch("integrations.aws.identity_store.execute_aws_api_call") +@patch("integrations.aws.identity_store.resolve_identity_store_id") +def test_delete_user(mock_resolve_identity_store_id, mock_execute_aws_api_call): + mock_resolve_identity_store_id.return_value = { + "IdentityStoreId": "test_instance_id" + } + mock_execute_aws_api_call.return_value = {} + user_id = "test_user_id" + + result = identity_store.delete_user(user_id) + + mock_execute_aws_api_call.assert_has_calls( + [ + call( + "identitystore", + "delete_user", + IdentityStoreId="test_instance_id", + UserId=user_id, + ) + ] + ) + assert result is True + + +@patch("integrations.aws.identity_store.execute_aws_api_call") +@patch("integrations.aws.identity_store.resolve_identity_store_id") +def test_delete_user_not_found( + mock_resolve_identity_store_id, mock_execute_aws_api_call +): + mock_resolve_identity_store_id.return_value = { + "IdentityStoreId": "test_instance_id" + } + mock_execute_aws_api_call.return_value = False + user_id = "nonexistent_user_id" + + result = identity_store.delete_user(user_id) + + mock_execute_aws_api_call.assert_has_calls( + [ + call( + "identitystore", + "delete_user", + IdentityStoreId="test_instance_id", + UserId=user_id, + ) + ] + ) + assert result is False + + @patch.dict(os.environ, {"AWS_SSO_INSTANCE_ID": "test_instance_id"}) @patch("integrations.utils.api.convert_string_to_camel_case") @patch("integrations.aws.identity_store.execute_aws_api_call")