fix: Add a fake arn to Context (#23)

* fix: Add a fake arn to Context * Fix DEFUALT_1P_ENTRYPOINT to DEFAULT_1P_ENTRYPOINT in tests Co-authored-by: Jacob Fuss <[email protected]>
aws · Mar 24, 2021 · 00d793b · 00d793b
1 parent 2c9ea44
commit 00d793b
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 15 deletions.
diff --git a/cmd/aws-lambda-rie/handlers.go b/cmd/aws-lambda-rie/handlers.go
@@ -94,11 +94,12 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox) {
 
 	invokeStart := time.Now()
 	invokePayload := &interop.Invoke{
-		ID:              uuid.New().String(),
-		TraceID:         r.Header.Get("X-Amzn-Trace-Id"),
-		LambdaSegmentID: r.Header.Get("X-Amzn-Segment-Id"),
-		Payload:         bodyBytes,
-		CorrelationID:   "invokeCorrelationID",
+		ID:                 uuid.New().String(),
+		InvokedFunctionArn: fmt.Sprintf("arn:aws:lambda:us-east-1:012345678912:function:%s", GetenvWithDefault("AWS_LAMBDA_FUNCTION_NAME", "test_function")),
+		TraceID:            r.Header.Get("X-Amzn-Trace-Id"),
+		LambdaSegmentID:    r.Header.Get("X-Amzn-Segment-Id"),
+		Payload:            bodyBytes,
+		CorrelationID:      "invokeCorrelationID",
 	}
 	fmt.Println("START RequestId: " + invokePayload.ID + " Version: " + functionVersion)
 
@@ -176,7 +177,6 @@ func InitHandler(sandbox Sandbox, functionVersion string, timeout int64) (time.T
 	additionalFunctionEnvironmentVariables["AWS_LAMBDA_FUNCTION_MEMORY_SIZE"] = "3008"
 	additionalFunctionEnvironmentVariables["AWS_LAMBDA_FUNCTION_NAME"] = "test_function"
 
-
 	// Forward Env Vars from the running system (container) to what the function can view. Without this, Env Vars will
 	// not be viewable when the function runs.
 	for _, env := range os.Environ() {
@@ -185,8 +185,6 @@ func InitHandler(sandbox Sandbox, functionVersion string, timeout int64) (time.T
 		additionalFunctionEnvironmentVariables[envVar[0]] = envVar[1]
 	}
 
-
-
 	initStart := time.Now()
 	// pass to rapid
 	sandbox.Init(&interop.Init{

diff --git a/test/integration/local_lambda/end-to-end-test.py b/test/integration/local_lambda/end-to-end-test.py
@@ -9,7 +9,7 @@
 import requests
 
 SLEEP_TIME = 2
-DEFUALT_1P_ENTRYPOINT = "/lambda-entrypoint.sh"
+DEFAULT_1P_ENTRYPOINT = "/lambda-entrypoint.sh"
 
 class TestEndToEnd(TestCase):
 
@@ -36,7 +36,7 @@ def tearDownClass(cls):
 
 
     def test_env_var_with_eqaul_sign(self):
-        cmd = f"docker run --name envvarcheck -d -v {self.path_to_binary}:/local-lambda-runtime-server -p 9003:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFUALT_1P_ENTRYPOINT} main.check_env_var_handler"
+        cmd = f"docker run --name envvarcheck -d -v {self.path_to_binary}:/local-lambda-runtime-server -p 9003:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFAULT_1P_ENTRYPOINT} main.check_env_var_handler"
 
         Popen(cmd.split(' ')).communicate()
 
@@ -47,7 +47,7 @@ def test_env_var_with_eqaul_sign(self):
         self.assertEqual(b'"4=4"', r.content)
 
     def test_two_invokes(self):
-        cmd = f"docker run --name testing -d -v {self.path_to_binary}:/local-lambda-runtime-server -p 9000:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFUALT_1P_ENTRYPOINT} main.success_handler"
+        cmd = f"docker run --name testing -d -v {self.path_to_binary}:/local-lambda-runtime-server -p 9000:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFAULT_1P_ENTRYPOINT} main.success_handler"
 
         Popen(cmd.split(' ')).communicate()
 
@@ -61,9 +61,31 @@ def test_two_invokes(self):
         r = requests.post("http://localhost:9000/2015-03-31/functions/function/invocations", json={})
         self.assertEqual(b'"My lambda ran succesfully"', r.content) 
 
+    def test_lambda_function_arn_exists(self):
+        cmd = f"docker run --name testing -d -v {self.path_to_binary}:/local-lambda-runtime-server -p 9000:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFAULT_1P_ENTRYPOINT} main.assert_lambda_arn_in_context"
+
+        Popen(cmd.split(' ')).communicate()
+
+        # sleep 1s to give enough time for the endpoint to be up to curl
+        time.sleep(SLEEP_TIME)
+
+        r = requests.post("http://localhost:9000/2015-03-31/functions/function/invocations", json={})
+        self.assertEqual(b'"My lambda ran succesfully"', r.content)
+
+    def test_lambda_function_arn_exists_with_defining_custom_name(self):
+        cmd = f"docker run --name testing --env AWS_LAMBDA_FUNCTION_NAME=MyCoolName -d -v {self.path_to_binary}:/local-lambda-runtime-server -p 9000:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFAULT_1P_ENTRYPOINT} main.assert_lambda_arn_in_context"
+
+        Popen(cmd.split(' ')).communicate()
+
+        # sleep 1s to give enough time for the endpoint to be up to curl
+        time.sleep(SLEEP_TIME)
+
+        r = requests.post("http://localhost:9000/2015-03-31/functions/function/invocations", json={})
+        self.assertEqual(b'"My lambda ran succesfully"', r.content)
+
 
     def test_timeout_invoke(self):
-        cmd = f"docker run --name timeout -d --env AWS_LAMBDA_FUNCTION_TIMEOUT=1 -v {self.path_to_binary}:/local-lambda-runtime-server -p 9001:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFUALT_1P_ENTRYPOINT} main.sleep_handler"
+        cmd = f"docker run --name timeout -d --env AWS_LAMBDA_FUNCTION_TIMEOUT=1 -v {self.path_to_binary}:/local-lambda-runtime-server -p 9001:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFAULT_1P_ENTRYPOINT} main.sleep_handler"
 
         Popen(cmd.split(' ')).communicate()
 
@@ -74,7 +96,7 @@ def test_timeout_invoke(self):
         self.assertEqual(b"Task timed out after 1.00 seconds", r.content)
 
     def test_exception_returned(self):
-        cmd = f"docker run --name exception -d -v {self.path_to_binary}:/local-lambda-runtime-server -p 9002:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFUALT_1P_ENTRYPOINT} main.exception_handler"
+        cmd = f"docker run --name exception -d -v {self.path_to_binary}:/local-lambda-runtime-server -p 9002:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFAULT_1P_ENTRYPOINT} main.exception_handler"
 
         Popen(cmd.split(' ')).communicate()
 
@@ -108,7 +130,7 @@ def tearDownClass(cls):
         Popen(f"docker rmi {cls.image_name}".split(' ')).communicate()
 
     def test_invoke_with_pre_runtime_api_runtime(self):
-        cmd = f"docker run --name testing -d -v {self.path_to_binary}:/local-lambda-runtime-server -p 9000:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFUALT_1P_ENTRYPOINT} main.success_handler"
+        cmd = f"docker run --name testing -d -v {self.path_to_binary}:/local-lambda-runtime-server -p 9000:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFAULT_1P_ENTRYPOINT} main.success_handler"
 
         Popen(cmd.split(' ')).communicate()
 
@@ -119,7 +141,7 @@ def test_invoke_with_pre_runtime_api_runtime(self):
         self.assertEqual(b'"My lambda ran succesfully"', r.content)
 
     def test_function_name_is_overriden(self):
-        cmd = f"docker run --name assert-overwritten -d --env AWS_LAMBDA_FUNCTION_NAME=MyCoolName -v {self.path_to_binary}:/local-lambda-runtime-server -p 9009:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFUALT_1P_ENTRYPOINT} main.assert_env_var_is_overwritten"
+        cmd = f"docker run --name assert-overwritten -d --env AWS_LAMBDA_FUNCTION_NAME=MyCoolName -v {self.path_to_binary}:/local-lambda-runtime-server -p 9009:8080 --entrypoint /local-lambda-runtime-server/aws-lambda-rie {self.image_name} {DEFAULT_1P_ENTRYPOINT} main.assert_env_var_is_overwritten"
 
         Popen(cmd.split(' ')).communicate()
 

diff --git a/test/integration/testdata/main.py b/test/integration/testdata/main.py
@@ -29,3 +29,9 @@ def assert_env_var_is_overwritten(event, context):
         raise("Function name was not overwritten")
     else:
         return "My lambda ran succesfully"
+
+def assert_lambda_arn_in_context(event, context):
+    if context.invoked_function_arn == f"arn:aws:lambda:us-east-1:012345678912:function:{os.environ.get('AWS_LAMBDA_FUNCTION_NAME', 'test_function')}":
+        return "My lambda ran succesfully"
+    else:
+        raise("Function Arn was not there")