diff --git a/app/integrations/aws/client.py b/app/integrations/aws/client.py index 83e07ec6..66ba5b77 100644 --- a/app/integrations/aws/client.py +++ b/app/integrations/aws/client.py @@ -12,9 +12,11 @@ SYSTEM_ADMIN_PERMISSIONS = os.environ.get("AWS_SSO_SYSTEM_ADMIN_PERMISSIONS") VIEW_ONLY_PERMISSIONS = os.environ.get("AWS_SSO_VIEW_ONLY_PERMISSIONS") AWS_REGION = os.environ.get("AWS_REGION", "ca-central-1") +THROTTLING_ERRORS = ["Throttling", "ThrottlingException", "RequestLimitExceeded"] +RESOURCE_NOT_FOUND_ERRORS = ["ResourceNotFoundException", "NoSuchEntity"] -logger = logging.getLogger(__name__) +logger = logging.getLogger() def handle_aws_api_errors(func): @@ -32,22 +34,17 @@ def wrapper(*args, **kwargs): try: return func(*args, **kwargs) except BotoCoreError as e: - logger.error( - f"A BotoCore error occurred in function '{func.__name__}': {e}" - ) + logger.error(f"{func.__module__}.{func.__name__}:BotoCore error: {e}") except ClientError as e: - if e.response["Error"]["Code"] == "ResourceNotFoundException": - logger.info(f"Resource not found in function '{func.__name__}': {e}") - return False + if e.response["Error"]["Code"] in THROTTLING_ERRORS: + logger.info(f"{func.__module__}.{func.__name__}: {e}") + elif e.response["Error"]["Code"] in RESOURCE_NOT_FOUND_ERRORS: + logger.warn(f"{func.__module__}.{func.__name__}: {e}") else: - logger.error( - f"A ClientError occurred in function '{func.__name__}': {e}" - ) + logger.error(f"{func.__module__}.{func.__name__}: {e}") except Exception as e: # Catch-all for any other types of exceptions - logger.error( - f"An unexpected error occurred in function '{func.__name__}': {e}" - ) - return None + logger.error(f"{func.__module__}.{func.__name__}: {e}") + return False return wrapper diff --git a/app/integrations/aws/identity_store.py b/app/integrations/aws/identity_store.py index 9f1c3972..be493e2d 100644 --- a/app/integrations/aws/identity_store.py +++ b/app/integrations/aws/identity_store.py @@ -63,7 +63,8 @@ def delete_user(user_id, **kwargs): kwargs = resolve_identity_store_id(kwargs) kwargs.update({"UserId": user_id}) response = execute_aws_api_call("identitystore", "delete_user", **kwargs) - return True if response == {} else False + del response["ResponseMetadata"] + return response == {} @handle_aws_api_errors @@ -85,8 +86,7 @@ def get_user_id(user_name, **kwargs): } } ) - response = execute_aws_api_call("identitystore", "get_user_id", **kwargs) - return response["UserId"] if response else False + return execute_aws_api_call("identitystore", "get_user_id", **kwargs)["UserId"] @handle_aws_api_errors @@ -100,9 +100,7 @@ def describe_user(user_id, **kwargs): kwargs = resolve_identity_store_id(kwargs) kwargs.update({"UserId": user_id}) response = execute_aws_api_call("identitystore", "describe_user", **kwargs) - if not response: - return False - response.pop("ResponseMetadata", None) + del response["ResponseMetadata"] return response @@ -122,6 +120,9 @@ def get_group_id(group_name, **kwargs): Args: group_name (str): The name of the group. **kwargs: Additional keyword arguments for the API call. + + Returns: + str: The group ID of the group. """ kwargs = resolve_identity_store_id(kwargs) kwargs.update( @@ -140,7 +141,13 @@ def get_group_id(group_name, **kwargs): @handle_aws_api_errors def list_groups(**kwargs): - """Retrieves all groups from the AWS Identity Center (identitystore)""" + """Retrieves all groups from the AWS Identity Center (identitystore) + + Args: + **kwargs: Additional keyword arguments for the API call. + + Returns: + list: A list of group objects.""" kwargs = resolve_identity_store_id(kwargs) response = execute_aws_api_call( "identitystore", "list_groups", paginated=True, keys=["Groups"], **kwargs @@ -184,7 +191,8 @@ def delete_group_membership(membership_id, **kwargs): response = execute_aws_api_call( "identitystore", "delete_group_membership", **kwargs ) - return True if response["ResponseMetadata"]["HTTPStatusCode"] == 200 else False + del response["ResponseMetadata"] + return response == {} @handle_aws_api_errors diff --git a/app/tests/integrations/aws/test_client.py b/app/tests/integrations/aws/test_client.py index fe35b75f..5876d76b 100644 --- a/app/tests/integrations/aws/test_client.py +++ b/app/tests/integrations/aws/test_client.py @@ -7,86 +7,79 @@ ROLE_ARN = "test_role_arn" -@patch("integrations.aws.client.logger.error") -@patch("integrations.aws.client.logger.info") -def test_handle_aws_api_errors_catches_botocore_error( - mocked_logging_info, mocked_logging_error -): +@patch("integrations.aws.client.logger") +def test_handle_aws_api_errors_catches_botocore_error(mock_logger): mock_func = MagicMock(side_effect=BotoCoreError()) - mock_func.__name__ = "mock_func" + mock_func.__name__ = "mock_func_name" + mock_func.__module__ = "mock_module" decorated_func = aws_client.handle_aws_api_errors(mock_func) result = decorated_func() - assert result is None + assert result is False mock_func.assert_called_once() - mocked_logging_error.assert_called_once_with( - "A BotoCore error occurred in function 'mock_func': An unspecified error occurred" + mock_logger.error.assert_called_once_with( + "mock_module.mock_func_name:BotoCore error: An unspecified error occurred" ) - mocked_logging_info.assert_not_called() + mock_logger.info.assert_not_called() -@patch("integrations.aws.client.logger.error") -@patch("integrations.aws.client.logger.info") -def test_handle_aws_api_errors_catches_client_error_resource_not_found( - mocked_logging_info, mocked_logging_error -): +@patch("integrations.aws.client.logger") +def test_handle_aws_api_errors_catches_client_error_resource_not_found(mock_logger): mock_func = MagicMock( side_effect=ClientError( {"Error": {"Code": "ResourceNotFoundException"}}, "operation_name" ) ) - mock_func.__name__ = "mock_func" + mock_func.__name__ = "mock_func_name" + mock_func.__module__ = "mock_module" decorated_func = aws_client.handle_aws_api_errors(mock_func) result = decorated_func() assert result is False mock_func.assert_called_once() - mocked_logging_info.assert_called_once_with( - "Resource not found in function 'mock_func': An error occurred (ResourceNotFoundException) when calling the operation_name operation: Unknown" + mock_logger.warn.assert_called_once_with( + "mock_module.mock_func_name: An error occurred (ResourceNotFoundException) when calling the operation_name operation: Unknown" ) - mocked_logging_error.assert_not_called() + mock_logger.error.assert_not_called() + mock_logger.info.assert_not_called() -@patch("integrations.aws.client.logger.error") -@patch("integrations.aws.client.logger.info") -def test_handle_aws_api_errors_catches_client_error_other( - mocked_logging_info, mocked_logging_error -): +@patch("integrations.aws.client.logger") +def test_handle_aws_api_errors_catches_client_error_other(mock_logger): mock_func = MagicMock( side_effect=ClientError({"Error": {"Code": "OtherError"}}, "operation_name") ) - mock_func.__name__ = "mock_func" + mock_func.__name__ = "mock_func_name" + mock_func.__module__ = "mock_module" decorated_func = aws_client.handle_aws_api_errors(mock_func) result = decorated_func() - assert result is None + assert result is False mock_func.assert_called_once() - mocked_logging_error.assert_called_once_with( - "A ClientError occurred in function 'mock_func': An error occurred (OtherError) when calling the operation_name operation: Unknown" + mock_logger.error.assert_called_once_with( + "mock_module.mock_func_name: An error occurred (OtherError) when calling the operation_name operation: Unknown" ) - mocked_logging_info.assert_not_called() + mock_logger.info.assert_not_called() -@patch("integrations.aws.client.logger.error") -@patch("integrations.aws.client.logger.info") -def test_handle_aws_api_errors_catches_exception( - mocked_logging_info, mocked_logging_error -): +@patch("integrations.aws.client.logger") +def test_handle_aws_api_errors_catches_exception(mock_logger): mock_func = MagicMock(side_effect=Exception("Exception message")) - mock_func.__name__ = "mock_func" + mock_func.__name__ = "mock_func_name" + mock_func.__module__ = "mock_module" decorated_func = aws_client.handle_aws_api_errors(mock_func) result = decorated_func() - assert result is None + assert result is False mock_func.assert_called_once() - mocked_logging_error.assert_called_once_with( - "An unexpected error occurred in function 'mock_func': Exception message" + mock_logger.error.assert_called_once_with( + "mock_module.mock_func_name: Exception message" ) - mocked_logging_info.assert_not_called() + mock_logger.info.assert_not_called() def test_handle_aws_api_errors_passes_through_return_value(): diff --git a/app/tests/integrations/aws/test_identity_store.py b/app/tests/integrations/aws/test_identity_store.py index f5bc23a9..329c3f71 100644 --- a/app/tests/integrations/aws/test_identity_store.py +++ b/app/tests/integrations/aws/test_identity_store.py @@ -106,7 +106,7 @@ def test_create_user_failed(mock_resolve_identity_store_id, mock_execute_aws_api Name={"GivenName": first_name, "FamilyName": family_name}, DisplayName=f"{first_name} {family_name}", ) - assert result is None + assert result is False @patch("integrations.aws.identity_store.execute_aws_api_call") @@ -179,6 +179,7 @@ def test_describe_user( mock_resolve_identity_store_id.return_value = { "IdentityStoreId": "test_instance_id" } + user["ResponseMetadata"] = {"HTTPStatusCode": 200} mock_execute_aws_api_call.return_value = user user_id = "user_id1" @@ -239,7 +240,9 @@ 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 = {} + mock_execute_aws_api_call.return_value = { + "ResponseMetadata": {"HTTPStatusCode": 200} + } user_id = "test_user_id" result = identity_store.delete_user(user_id) @@ -593,9 +596,8 @@ def test_delete_group_membership_resource_not_found( mock_resolve_identity_store_id.return_value = { "IdentityStoreId": "test_instance_id" } - mock_execute_aws_api_call.return_value = { - "ResponseMetadata": {"HTTPStatusCode": 404} - } + # API error handling should return False if the resource is not found + mock_execute_aws_api_call.return_value = False membership_id = "nonexistent_membership_id" result = identity_store.delete_group_membership(membership_id)