fix: handling of AWS API responses (#497)

* fix: handling of AWS API responses

* fix: lint and fmt

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
 

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

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():

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)