diff --git a/moto/awslambda/models.py b/moto/awslambda/models.py index 13d4726ac..523393ba3 100644 --- a/moto/awslambda/models.py +++ b/moto/awslambda/models.py @@ -81,6 +81,8 @@ class LambdaFunction(BaseModel): self.function_arn = 'arn:aws:lambda:123456789012:function:{0}'.format( self.function_name) + self.tags = dict() + @property def vpc_config(self): config = self._vpc_config.copy() @@ -278,6 +280,9 @@ class LambdaBackend(BaseBackend): def has_function(self, function_name): return function_name in self._functions + def has_function_arn(self, function_arn): + return self.get_function_by_arn(function_arn) is not None + def create_function(self, spec): fn = LambdaFunction(spec) self._functions[fn.function_name] = fn @@ -286,12 +291,33 @@ class LambdaBackend(BaseBackend): def get_function(self, function_name): return self._functions[function_name] + def get_function_by_arn(self, function_arn): + for function in self._functions.values(): + if function.function_arn == function_arn: + return function + return None + def delete_function(self, function_name): del self._functions[function_name] def list_functions(self): return self._functions.values() + def list_tags(self, resource): + return self.get_function_by_arn(resource).tags + + def tag_resource(self, resource, tags): + self.get_function_by_arn(resource).tags.update(tags) + + def untag_resource(self, resource, tagKeys): + function = self.get_function_by_arn(resource) + for key in tagKeys: + try: + del function.tags[key] + except KeyError: + pass + # Don't care + def do_validate_s3(): return os.environ.get('VALIDATE_LAMBDA_S3', '') in ['', '1', 'true'] diff --git a/moto/awslambda/responses.py b/moto/awslambda/responses.py index d145f4760..b7a6921c5 100644 --- a/moto/awslambda/responses.py +++ b/moto/awslambda/responses.py @@ -3,6 +3,12 @@ from __future__ import unicode_literals import json import re +try: + from urllib import unquote + from urlparse import urlparse, parse_qs +except: + from urllib.parse import unquote, urlparse, parse_qs + from moto.core.responses import BaseResponse @@ -33,6 +39,17 @@ class LambdaResponse(BaseResponse): else: raise ValueError("Cannot handle request") + def tag(self, request, full_url, headers): + self.setup_class(request, full_url, headers) + if request.method == 'GET': + return self._list_tags(request, full_url) + elif request.method == 'POST': + return self._tag_resource(request, full_url) + elif request.method == 'DELETE': + return self._untag_resource(request, full_url) + else: + raise ValueError("Cannot handle {0} request".format(request.method)) + def _invoke(self, request, full_url): response_headers = {} lambda_backend = self.get_lambda_backend(full_url) @@ -102,3 +119,43 @@ class LambdaResponse(BaseResponse): return region.group(1) else: return self.default_region + + def _list_tags(self, request, full_url): + lambda_backend = self.get_lambda_backend(full_url) + + path = request.path if hasattr(request, 'path') else request.path_url + function_arn = unquote(path.split('/')[-1]) + + if lambda_backend.has_function_arn(function_arn): + function = lambda_backend.get_function_by_arn(function_arn) + return 200, {}, json.dumps(dict(Tags=function.tags)) + else: + return 404, {}, "{}" + + def _tag_resource(self, request, full_url): + lambda_backend = self.get_lambda_backend(full_url) + + path = request.path if hasattr(request, 'path') else request.path_url + function_arn = unquote(path.split('/')[-1]) + + spec = json.loads(self.body) + + if lambda_backend.has_function_arn(function_arn): + lambda_backend.tag_resource(function_arn, spec['Tags']) + return 200, {}, "{}" + else: + return 404, {}, "{}" + + def _untag_resource(self, request, full_url): + lambda_backend = self.get_lambda_backend(full_url) + + path = request.path if hasattr(request, 'path') else request.path_url + function_arn = unquote(path.split('/')[-1].split('?')[0]) + + tag_keys = parse_qs(urlparse(full_url).query)['tagKeys'] + + if lambda_backend.has_function_arn(function_arn): + lambda_backend.untag_resource(function_arn, tag_keys) + return 204, {}, "{}" + else: + return 404, {}, "{}" diff --git a/moto/awslambda/urls.py b/moto/awslambda/urls.py index c63135766..c0917be57 100644 --- a/moto/awslambda/urls.py +++ b/moto/awslambda/urls.py @@ -11,4 +11,5 @@ url_paths = { '{0}/(?P[^/]+)/functions/?$': response.root, '{0}/(?P[^/]+)/functions/(?P[\w_-]+)/?$': response.function, '{0}/(?P[^/]+)/functions/(?P[\w_-]+)/invocations/?$': response.invoke, + '{0}/(?P[^/]+)/tags/(?P.+)': response.tag } diff --git a/tests/test_awslambda/test_lambda.py b/tests/test_awslambda/test_lambda.py index 007516f56..4ec504365 100644 --- a/tests/test_awslambda/test_lambda.py +++ b/tests/test_awslambda/test_lambda.py @@ -469,3 +469,89 @@ def test_invoke_lambda_error(): assert 'FunctionError' in result assert result['FunctionError'] == 'Handled' + +@mock_lambda +@mock_s3 +def test_tags(): + """ + test list_tags -> tag_resource -> list_tags -> tag_resource -> list_tags -> untag_resource -> list_tags integration + """ + s3_conn = boto3.client('s3', 'us-west-2') + s3_conn.create_bucket(Bucket='test-bucket') + + zip_content = get_test_zip_file2() + s3_conn.put_object(Bucket='test-bucket', Key='test.zip', Body=zip_content) + conn = boto3.client('lambda', 'us-west-2') + + function = conn.create_function( + FunctionName='testFunction', + Runtime='python2.7', + Role='test-iam-role', + Handler='lambda_function.handler', + Code={ + 'S3Bucket': 'test-bucket', + 'S3Key': 'test.zip', + }, + Description='test lambda function', + Timeout=3, + MemorySize=128, + Publish=True, + ) + + # List tags when there are none + conn.list_tags( + Resource=function['FunctionArn'] + )['Tags'].should.equal(dict()) + + # List tags when there is one + conn.tag_resource( + Resource=function['FunctionArn'], + Tags=dict(spam='eggs') + )['ResponseMetadata']['HTTPStatusCode'].should.equal(200) + conn.list_tags( + Resource=function['FunctionArn'] + )['Tags'].should.equal(dict(spam='eggs')) + + # List tags when another has been added + conn.tag_resource( + Resource=function['FunctionArn'], + Tags=dict(foo='bar') + )['ResponseMetadata']['HTTPStatusCode'].should.equal(200) + conn.list_tags( + Resource=function['FunctionArn'] + )['Tags'].should.equal(dict(spam='eggs', foo='bar')) + + # Untag resource + conn.untag_resource( + Resource=function['FunctionArn'], + TagKeys=['spam', 'trolls'] + )['ResponseMetadata']['HTTPStatusCode'].should.equal(204) + conn.list_tags( + Resource=function['FunctionArn'] + )['Tags'].should.equal(dict(foo='bar')) + + # Untag a tag that does not exist (no error and no change) + conn.untag_resource( + Resource=function['FunctionArn'], + TagKeys=['spam'] + )['ResponseMetadata']['HTTPStatusCode'].should.equal(204) + +@mock_lambda +def test_tags_not_found(): + """ + Test list_tags and tag_resource when the lambda with the given arn does not exist + """ + conn = boto3.client('lambda', 'us-west-2') + conn.list_tags.when.called_with( + Resource='arn:aws:lambda:123456789012:function:not-found' + ).should.throw(botocore.client.ClientError) + + conn.tag_resource.when.called_with( + Resource='arn:aws:lambda:123456789012:function:not-found', + Tags=dict(spam='eggs') + ).should.throw(botocore.client.ClientError) + + conn.untag_resource.when.called_with( + Resource='arn:aws:lambda:123456789012:function:not-found', + TagKeys=['spam'] + ).should.throw(botocore.client.ClientError)