AWSLambda - Error when removing unknown permission (#4846)

moto/awslambda/exceptions.py

@@ -1,11 +1,9 @@
-from botocore.client import ClientError
 from moto.core.exceptions import JsonRESTError
 
 
-class LambdaClientError(ClientError):
+class LambdaClientError(JsonRESTError):
     def __init__(self, error, message):
-        error_response = {"Error": {"Code": error, "Message": message}}
-        super().__init__(error_response, None)
+        super().__init__(error, message)
 
 
 class CrossAccountNotAllowed(LambdaClientError):
@@ -35,3 +33,13 @@ class PreconditionFailedException(JsonRESTError):
 
     def __init__(self, message):
         super().__init__("PreconditionFailedException", message)
+
+
+class UnknownPolicyException(LambdaClientError):
+    code = 404
+
+    def __init__(self):
+        super().__init__(
+            "ResourceNotFoundException",
+            "No policy is associated with the given resource.",
+        )

moto/awslambda/policy.py

@@ -1,7 +1,10 @@
 import json
 import uuid
 
-from moto.awslambda.exceptions import PreconditionFailedException
+from moto.awslambda.exceptions import (
+    PreconditionFailedException,
+    UnknownPolicyException,
+)
 
 
 class Policy:
@@ -48,6 +51,9 @@ class Policy:
         for statement in self.statements:
             if "Sid" in statement and statement["Sid"] == sid:
                 self.statements.remove(statement)
+                break
+        else:
+            raise UnknownPolicyException()
 
     # converts AddPermission request to PolicyStatement
     # https://docs.aws.amazon.com/lambda/latest/dg/API_AddPermission.html

moto/awslambda/responses.py

@@ -6,11 +6,24 @@ try:
 except ImportError:
     from urllib.parse import unquote
 
+from functools import wraps
 from moto.core.utils import amz_crc32, amzn_request_id, path_url
 from moto.core.responses import BaseResponse
+from .exceptions import LambdaClientError
 from .models import lambda_backends
 
 
+def error_handler(f):
+    @wraps(f)
+    def _wrapper(*args, **kwargs):
+        try:
+            return f(*args, **kwargs)
+        except LambdaClientError as e:
+            return e.code, e.get_headers(), e.get_body()
+
+    return _wrapper
+
+
 class LambdaResponse(BaseResponse):
     @property
     def json_body(self):
@@ -29,6 +42,7 @@ class LambdaResponse(BaseResponse):
         """
         return lambda_backends[self.region]
 
+    @error_handler
     def root(self, request, full_url, headers):
         self.setup_class(request, full_url, headers)
         if request.method == "GET":
@@ -127,6 +141,7 @@ class LambdaResponse(BaseResponse):
         else:
             raise ValueError("Cannot handle {0} request".format(request.method))
 
+    @error_handler
     def policy(self, request, full_url, headers):
         self.setup_class(request, full_url, headers)
         if request.method == "GET":

tests/test_awslambda/test_lambda.py

@@ -5,6 +5,7 @@ import json
 import sure  # noqa # pylint: disable=unused-import
 import pytest
 
+from botocore.exceptions import ClientError
 from freezegun import freeze_time
 from moto import mock_lambda, mock_s3
 from moto.core.models import ACCOUNT_ID
@@ -1227,3 +1228,24 @@ def test_remove_function_permission(key):
     policy = conn.get_policy(FunctionName=name_or_arn, Qualifier="2")["Policy"]
     policy = json.loads(policy)
     policy["Statement"].should.equal([])
+
+
+@mock_lambda
+def test_remove_unknown_permission_throws_error():
+    conn = boto3.client("lambda", _lambda_region)
+    zip_content = get_test_zip_file1()
+    function_name = str(uuid4())[0:6]
+    f = conn.create_function(
+        FunctionName=function_name,
+        Runtime="python3.7",
+        Role=(get_role_name()),
+        Handler="lambda_function.handler",
+        Code={"ZipFile": zip_content},
+    )
+    arn = f["FunctionArn"]
+
+    with pytest.raises(ClientError) as exc:
+        conn.remove_permission(FunctionName=arn, StatementId="1")
+    err = exc.value.response["Error"]
+    err["Code"].should.equal("ResourceNotFoundException")
+    err["Message"].should.equal("No policy is associated with the given resource.")