From 941d817da4e62c72f638e4972e3ae39dfbcf4d7b Mon Sep 17 00:00:00 2001
From: Alexander Mohr <thehesiod@users.noreply.github.com>
Date: Wed, 21 Mar 2018 22:14:10 -0700
Subject: [PATCH] fix lambda endpoint parsing (#1395)

* fix endpoint parsing

* add new unittest

* finish test
---
 moto/awslambda/models.py            |  6 +--
 moto/sns/models.py                  | 18 ++++++--
 tests/test_awslambda/test_lambda.py | 65 +++++++++++++++++++++++++++--
 3 files changed, 80 insertions(+), 9 deletions(-)

diff --git a/moto/awslambda/models.py b/moto/awslambda/models.py
index 3c3d3ea66..80b4ffba3 100644
--- a/moto/awslambda/models.py
+++ b/moto/awslambda/models.py
@@ -603,7 +603,7 @@ class LambdaBackend(BaseBackend):
     def list_functions(self):
         return self._lambdas.all()
 
-    def send_message(self, function_name, message, subject=None):
+    def send_message(self, function_name, message, subject=None, qualifier=None):
         event = {
             "Records": [
                 {
@@ -636,8 +636,8 @@ class LambdaBackend(BaseBackend):
             ]
 
         }
-        self._functions[function_name][-1].invoke(json.dumps(event), {}, {})
-        pass
+        func = self._lambdas.get_function(function_name, qualifier)
+        func.invoke(json.dumps(event), {}, {})
 
     def list_tags(self, resource):
         return self.get_function_by_arn(resource).tags
diff --git a/moto/sns/models.py b/moto/sns/models.py
index 9afc28f46..a66523614 100644
--- a/moto/sns/models.py
+++ b/moto/sns/models.py
@@ -100,9 +100,21 @@ class Subscription(BaseModel):
             requests.post(self.endpoint, json=post_data)
         elif self.protocol == 'lambda':
             # TODO: support bad function name
-            function_name = self.endpoint.split(":")[-1]
-            region = self.arn.split(':')[3]
-            lambda_backends[region].send_message(function_name, message, subject=subject)
+            # http://docs.aws.amazon.com/general/latest/gr/aws-arns-and-namespaces.html
+            arr = self.endpoint.split(":")
+            region = arr[3]
+            qualifier = None
+            if len(arr) == 7:
+                assert arr[5] == 'function'
+                function_name = arr[-1]
+            elif len(arr) == 8:
+                assert arr[5] == 'function'
+                qualifier = arr[-1]
+                function_name = arr[-2]
+            else:
+                assert False
+
+            lambda_backends[region].send_message(function_name, message, subject=subject, qualifier=qualifier)
 
     def _matches_filter_policy(self, message_attributes):
         # TODO: support Anything-but matching, prefix matching and
diff --git a/tests/test_awslambda/test_lambda.py b/tests/test_awslambda/test_lambda.py
index e7a9f9174..8ea9cc6fd 100644
--- a/tests/test_awslambda/test_lambda.py
+++ b/tests/test_awslambda/test_lambda.py
@@ -6,11 +6,12 @@ import boto3
 import hashlib
 import io
 import json
+import time
 import zipfile
 import sure  # noqa
 
 from freezegun import freeze_time
-from moto import mock_lambda, mock_s3, mock_ec2, settings
+from moto import mock_lambda, mock_s3, mock_ec2, mock_sns, mock_logs, settings
 
 _lambda_region = 'us-west-2'
 
@@ -48,6 +49,15 @@ def lambda_handler(event, context):
     return _process_lambda(func_str)
 
 
+def get_test_zip_file3():
+    pfunc = """
+def lambda_handler(event, context):
+    print("get_test_zip_file3 success")
+    return event
+"""
+    return _process_lambda(pfunc)
+
+
 @mock_lambda
 def test_list_functions():
     conn = boto3.client('lambda', 'us-west-2')
@@ -160,6 +170,56 @@ if settings.TEST_SERVER_MODE:
         payload.should.equal(msg)
 
 
+@mock_logs
+@mock_sns
+@mock_ec2
+@mock_lambda
+def test_invoke_function_from_sns():
+    logs_conn = boto3.client("logs", region_name="us-west-2")
+    sns_conn = boto3.client("sns", region_name="us-west-2")
+    sns_conn.create_topic(Name="some-topic")
+    topics_json = sns_conn.list_topics()
+    topics = topics_json["Topics"]
+    topic_arn = topics[0]['TopicArn']
+
+    conn = boto3.client('lambda', 'us-west-2')
+    result = conn.create_function(
+        FunctionName='testFunction',
+        Runtime='python2.7',
+        Role='test-iam-role',
+        Handler='lambda_function.lambda_handler',
+        Code={
+            'ZipFile': get_test_zip_file3(),
+        },
+        Description='test lambda function',
+        Timeout=3,
+        MemorySize=128,
+        Publish=True,
+    )
+
+    sns_conn.subscribe(TopicArn=topic_arn, Protocol="lambda", Endpoint=result['FunctionArn'])
+
+    result = sns_conn.publish(TopicArn=topic_arn, Message=json.dumps({}))
+
+    start = time.time()
+    while (time.time() - start) < 30:
+        result = logs_conn.describe_log_streams(logGroupName='/aws/lambda/testFunction')
+        log_streams = result.get('logStreams')
+        if not log_streams:
+            time.sleep(1)
+            continue
+
+        assert len(log_streams) == 1
+        result = logs_conn.get_log_events(logGroupName='/aws/lambda/testFunction', logStreamName=log_streams[0]['logStreamName'])
+        for event in result.get('events'):
+            if event['message'] == 'get_test_zip_file3 success':
+                return
+
+        time.sleep(1)
+
+    assert False, "Test Failed"
+
+
 @mock_lambda
 def test_create_based_on_s3_with_missing_bucket():
     conn = boto3.client('lambda', 'us-west-2')
@@ -420,7 +480,6 @@ def test_publish():
     function_list['Functions'][0]['FunctionArn'].should.contain('testFunction:$LATEST')
 
 
-
 @mock_lambda
 @mock_s3
 @freeze_time('2015-01-01 00:00:00')
@@ -674,7 +733,7 @@ def test_get_function_created_with_zipfile():
             "CodeSha256": hashlib.sha256(zip_content).hexdigest(),
             "CodeSize": len(zip_content),
             "Description": "test lambda function",
-            "FunctionArn":'arn:aws:lambda:{}:123456789012:function:testFunction:$LATEST'.format(_lambda_region),
+            "FunctionArn": 'arn:aws:lambda:{}:123456789012:function:testFunction:$LATEST'.format(_lambda_region),
             "FunctionName": "testFunction",
             "Handler": "lambda_function.handler",
             "MemorySize": 128,