refactor test cases

roger-zhangg · roger-zhangg · commit 2a1a2fed7e3c · 2024-04-11T17:59:09.000Z
diff --git a/Makefile b/Makefile
@@ -39,9 +39,32 @@ integ-tests-and-compile: tests
 integ-tests-with-docker: tests-with-docker
 	make compile-with-docker-all
 	make integ-tests
-	
-integ-tests:
+
+prep-python:
 	python3 -m venv .venv
 	.venv/bin/pip install --upgrade pip
 	.venv/bin/pip install requests parameterized
-	.venv/bin/python3 test/integration/local_lambda/test_end_to_end.py
+
+integ-tests:
+	make prep-python
+	docker run --rm --privileged multiarch/qemu-user-static --reset -p yes
+	TEST_ARCH=x86_64 TEST_PORT=8002 .venv/bin/python3 test/integration/local_lambda/test_end_to_end.py
+	TEST_ARCH=arm64 TEST_PORT=9002 .venv/bin/python3 test/integration/local_lambda/test_end_to_end.py
+	TEST_ARCH="" TEST_PORT=9052 .venv/bin/python3 test/integration/local_lambda/test_end_to_end.py
+
+integ-tests-with-docker-x86-64:
+	make ARCH=x86_64 compile-with-docker
+	make prep-python
+	TEST_ARCH=x86_64 TEST_PORT=8002 .venv/bin/python3 test/integration/local_lambda/test_end_to_end.py
+
+integ-tests-with-docker-arm64:
+	make ARCH=arm64 compile-with-docker
+	make prep-python
+	docker run --rm --privileged multiarch/qemu-user-static --reset -p yes
+	TEST_ARCH="arm64" TEST_PORT=9002 .venv/bin/python3 test/integration/local_lambda/test_end_to_end.py
+
+integ-tests-with-docker-old:
+	make ARCH=old compile-with-docker
+	make prep-python
+	TEST_ARCH="" TEST_PORT=9052 .venv/bin/python3 test/integration/local_lambda/test_end_to_end.py
+	
diff --git a/test/integration/local_lambda/test_end_to_end.py b/test/integration/local_lambda/test_end_to_end.py
@@ -5,16 +5,20 @@
 from unittest import TestCase, main
 from pathlib import Path
 import time
-
+import os
 import requests
+from contextlib import contextmanager
 from parameterized import parameterized
 
 SLEEP_TIME = 2
 DEFAULT_1P_ENTRYPOINT = "/lambda-entrypoint.sh"
 ARCHS = ["x86_64", "arm64", ""]
 
 
+
 class TestEndToEnd(TestCase):
+    ARCH = os.environ.get('TEST_ARCH', "")
+    PORT = os.environ.get('TEST_PORT', 8002)
     @classmethod
     def setUpClass(cls):
         testdata_path = Path(__file__).resolve().parents[1].joinpath("testdata")
@@ -23,48 +27,30 @@ def setUpClass(cls):
         cls.path_to_binary = Path().resolve().joinpath("bin")
 
         # build image
-        for arch in ARCHS:
-            image_name = cls.image_name if arch == "" else f"{cls.image_name}-{arch}"
-            architecture = arch if arch == "arm64" else "amd64"
-            build_cmd = [
-                "docker",
-                "build",
-                "--platform",
-                f"linux/{architecture}",
-                "-t",
-                image_name,
-                "-f",
-                str(dockerfile_path),
-                str(testdata_path),
-            ]
-            Popen(build_cmd).communicate()
+        image_name = cls.image_name if cls.ARCH == "" else f"{cls.image_name}-{cls.ARCH}"
+        architecture = cls.ARCH if cls.ARCH == "arm64" else "amd64"
+        docker_arch = cls.ARCH if cls.ARCH == "arm64" else "x86_64"
+        
+        build_cmd = [
+            "docker",
+            "build",
+            "--platform",
+            f"linux/{architecture}",
+            "-t",
+            image_name,
+            "-f",
+            str(dockerfile_path),
+            str(testdata_path),
+            "--build-arg",
+            f"IMAGE_ARCH={docker_arch}",
+        ]
+        print (build_cmd)
+        Popen(build_cmd).communicate()
 
     @classmethod
     def tearDownClass(cls):
-        images_to_delete = [
-            "envvarcheck",
-            "twoinvokes",
-            "arnexists",
-            "customname",
-            "timeout",
-            "exception",
-            "remaining_time_in_three_seconds",
-            "remaining_time_in_ten_seconds",
-            "remaining_time_in_default_deadline",
-            "pre-runtime-api",
-            "assert-overwritten",
-            "port_override"
-        ]
-
-        for image in images_to_delete:
-            for arch in ARCHS:
-                arch_tag = "" if arch == "" else f"-{arch}"
-                cmd = f"docker rm -f {image}{arch_tag}"
-                Popen(cmd.split(" ")).communicate()
-            
-        for arch in ARCHS:
-            arch_tag = "" if arch == "" else f"-{arch}"
-            Popen(f"docker rmi {cls.image_name}{arch_tag}".split(" ")).communicate()
+        arch_tag = "" if cls.ARCH == "" else f"-{cls.ARCH}"
+        Popen(f"docker rmi {cls.image_name}{arch_tag}".split(" ")).communicate()
 
     def tagged_name(self, name, architecture):
         tag = self.get_tag(architecture)
@@ -84,148 +70,170 @@ def invoke_function(self, port):
             f"http://localhost:{port}/2015-03-31/functions/function/invocations", json={}
         )
         
-    def create_container_and_invoke_function(self, cmd, port):
-        self.run_command(cmd)
-        
-        # sleep 1s to give enough time for the endpoint to be up to curl
-        self.sleep_1s()
-        
-        return self.invoke_function(port)
-
-    @parameterized.expand([("x86_64", "8000"), ("arm64", "9000"), ("", "9050")])
-    def test_env_var_with_equal_sign(self, arch, port):
+    @contextmanager
+    def create_container(self, cmd, image):
+        try:
+            platform = "x86_64" if self.ARCH == "" else self.ARCH
+            cmd_full = f"docker run --platform linux/{platform} {cmd}"
+            self.run_command(cmd_full)
+        
+            # sleep 1s to give enough time for the endpoint to be up to curl
+            self.sleep_1s()
+            yield
+        except Exception as e:
+            print(f"An error occurred while executing cmd: {cmd_full}. error: {e}")
+            raise e
+        finally:
+            self.run_command(f"docker stop {image}")
+            self.run_command(f"docker rm -f {image}")
+        
+
+    def test_env_var_with_equal_sign(self, arch=ARCH, port=PORT):
         image, rie, image_name = self.tagged_name("envvarcheck", arch)
-
-        cmd = f"docker run --name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.check_env_var_handler"
+        cmd = f"--name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.check_env_var_handler"
         
-        r = self.create_container_and_invoke_function(cmd, port)
+        with self.create_container(cmd, image):
+            r = self.invoke_function(port)
         
-        self.assertEqual(b'"4=4"', r.content)
+            self.assertEqual(b'"4=4"', r.content)
+
 
-    @parameterized.expand([("x86_64", "8001"), ("arm64", "9001"), ("", "9051")])
-    def test_two_invokes(self, arch, port):
+    def test_two_invokes(self, arch=ARCH, port=PORT):
         image, rie, image_name = self.tagged_name("twoinvokes", arch)
 
-        cmd = f"docker run --name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.success_handler"
+        cmd = f"--name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.success_handler"
 
-        r = self.create_container_and_invoke_function(cmd, port)
+        with self.create_container(cmd, image):
+            r = self.invoke_function(port)
         
-        self.assertEqual(b'"My lambda ran succesfully"', r.content)
+            self.assertEqual(b'"My lambda ran succesfully"', r.content)
+
+            # Make sure we can invoke the function twice
+            r = self.invoke_function(port)
+            
+            self.assertEqual(b'"My lambda ran succesfully"', r.content)
 
-        # Make sure we can invoke the function twice
-        r = self.invoke_function(port)
-        
-        self.assertEqual(b'"My lambda ran succesfully"', r.content)
 
-    @parameterized.expand([("x86_64", "8002"), ("arm64", "9002"), ("", "9052")])
-    def test_lambda_function_arn_exists(self, arch, port):
+    def test_lambda_function_arn_exists(self, arch=ARCH, port=PORT):
         image, rie, image_name = self.tagged_name("arnexists", arch)
 
-        cmd = f"docker run --name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.assert_lambda_arn_in_context"
+        cmd = f"--name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.assert_lambda_arn_in_context"
 
-        r = self.create_container_and_invoke_function(cmd, port)
+        with self.create_container(cmd, image):
+            r = self.invoke_function(port)
         
-        self.assertEqual(b'"My lambda ran succesfully"', r.content)
+            self.assertEqual(b'"My lambda ran succesfully"', r.content)
 
-    @parameterized.expand([("x86_64", "8003"), ("arm64", "9003"), ("", "9053")])
-    def test_lambda_function_arn_exists_with_defining_custom_name(self, arch, port):
+
+    def test_lambda_function_arn_exists_with_defining_custom_name(self, arch=ARCH, port=PORT):
         image, rie, image_name = self.tagged_name("customname", arch)
 
-        cmd = f"docker run --name {image} --env AWS_LAMBDA_FUNCTION_NAME=MyCoolName -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.assert_lambda_arn_in_context"
+        cmd = f"--name {image} --env AWS_LAMBDA_FUNCTION_NAME=MyCoolName -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.assert_lambda_arn_in_context"
         
-        r = self.create_container_and_invoke_function(cmd, port)
+        with self.create_container(cmd, image):
+            r = self.invoke_function(port)
         
-        self.assertEqual(b'"My lambda ran succesfully"', r.content)
+            self.assertEqual(b'"My lambda ran succesfully"', r.content)
+
 
-    @parameterized.expand([("x86_64", "8004"), ("arm64", "9004"), ("", "9054")])
-    def test_timeout_invoke(self, arch, port):
+    def test_timeout_invoke(self, arch=ARCH, port=PORT):
         image, rie, image_name = self.tagged_name("timeout", arch)
 
-        cmd = f"docker run --name {image} -d --env AWS_LAMBDA_FUNCTION_TIMEOUT=1 -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.sleep_handler"
+        cmd = f"--name {image} -d --env AWS_LAMBDA_FUNCTION_TIMEOUT=1 -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.sleep_handler"
 
-        r = self.create_container_and_invoke_function(cmd, port)
+        with self.create_container(cmd, image):
+            r = self.invoke_function(port)
         
-        self.assertEqual(b"Task timed out after 1.00 seconds", r.content)
+            self.assertEqual(b"Task timed out after 1.00 seconds", r.content)
 
-    @parameterized.expand([("x86_64", "8005"), ("arm64", "9005"), ("", "9055")])
-    def test_exception_returned(self, arch, port):
+
+    def test_exception_returned(self, arch=ARCH, port=PORT):
         image, rie, image_name = self.tagged_name("exception", arch)
 
-        cmd = f"docker run --name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.exception_handler"
+        cmd = f"--name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.exception_handler"
+
+        with self.create_container(cmd, image):
+            r = self.invoke_function(port)
+            
+            # ignore request_id in python3.12 lambda
+            result = r.json()
+            self.assertEqual(result["errorMessage"],"Raising an exception")
+            self.assertEqual(result["errorType"],"Exception")
+            self.assertEqual(result["stackTrace"],["  File \"/var/task/main.py\", line 13, in exception_handler\n    raise Exception(\"Raising an exception\")\n"])
 
-        r = self.create_container_and_invoke_function(cmd, port)
-        
-        self.assertEqual(
-            b'{"errorMessage": "Raising an exception", "errorType": "Exception", "stackTrace": ["  File \\"/var/task/main.py\\", line 13, in exception_handler\\n    raise Exception(\\"Raising an exception\\")\\n"]}',
-            r.content,
-        )
 
-    @parameterized.expand([("x86_64", "8006"), ("arm64", "9006"), ("", "9056")])
-    def test_context_get_remaining_time_in_three_seconds(self, arch, port):
+    def test_context_get_remaining_time_in_three_seconds(self, arch=ARCH, port=PORT):
         image, rie, image_name = self.tagged_name("remaining_time_in_three_seconds", arch)
 
-        cmd = f"docker run --name {image} -d --env AWS_LAMBDA_FUNCTION_TIMEOUT=3 -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.check_remaining_time_handler"
+        cmd = f"--name {image} -d --env AWS_LAMBDA_FUNCTION_TIMEOUT=3 -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.check_remaining_time_handler"
 
-        r = self.create_container_and_invoke_function(cmd, port)
+        with self.create_container(cmd, image):
+            r = self.invoke_function(port)
         
-        # Execution time is not decided, 1.0s ~ 3.0s is a good estimation
-        self.assertLess(int(r.content), 3000)
-        self.assertGreater(int(r.content), 1000)
+            # Execution time is not decided, 1.0s ~ 3.0s is a good estimation
+            self.assertLess(int(r.content), 3000)
+            self.assertGreater(int(r.content), 1000)
 
-    @parameterized.expand([("x86_64", "8007"), ("arm64", "9007"), ("", "9057")])
-    def test_context_get_remaining_time_in_ten_seconds(self, arch, port):
+
+    def test_context_get_remaining_time_in_ten_seconds(self, arch=ARCH, port=PORT):
         image, rie, image_name = self.tagged_name("remaining_time_in_ten_seconds", arch)
 
-        cmd = f"docker run --name {image} -d --env AWS_LAMBDA_FUNCTION_TIMEOUT=10 -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.check_remaining_time_handler"
+        cmd = f"--name {image} -d --env AWS_LAMBDA_FUNCTION_TIMEOUT=10 -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.check_remaining_time_handler"
 
-        r = self.create_container_and_invoke_function(cmd, port)
+        with self.create_container(cmd, image):
+            r = self.invoke_function(port)
         
-        # Execution time is not decided, 8.0s ~ 10.0s is a good estimation
-        self.assertLess(int(r.content), 10000)
-        self.assertGreater(int(r.content), 8000)
+            # Execution time is not decided, 8.0s ~ 10.0s is a good estimation
+            self.assertLess(int(r.content), 10000)
+            self.assertGreater(int(r.content), 8000)
+
 
-    @parameterized.expand([("x86_64", "8008"), ("arm64", "9008"), ("", "9058")])
-    def test_context_get_remaining_time_in_default_deadline(self, arch, port):
+    def test_context_get_remaining_time_in_default_deadline(self, arch=ARCH, port=PORT):
         image, rie, image_name = self.tagged_name("remaining_time_in_default_deadline", arch)
 
-        cmd = f"docker run --name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.check_remaining_time_handler"
+        cmd = f"--name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.check_remaining_time_handler"
+
+        with self.create_container(cmd, image):
+            r = self.invoke_function(port)
 
-        r = self.create_container_and_invoke_function(cmd, port)
+            # Executation time is not decided, 298.0s ~ 300.0s is a good estimation
+            self.assertLess(int(r.content), 300000)
+            self.assertGreater(int(r.content), 298000)
 
-        # Executation time is not decided, 298.0s ~ 300.0s is a good estimation
-        self.assertLess(int(r.content), 300000)
-        self.assertGreater(int(r.content), 298000)
 
-    @parameterized.expand([("x86_64", "8009"), ("arm64", "9009"), ("", "9059")])
-    def test_invoke_with_pre_runtime_api_runtime(self, arch, port):
+    def test_invoke_with_pre_runtime_api_runtime(self, arch=ARCH, port=PORT):
         image, rie, image_name = self.tagged_name("pre-runtime-api", arch)
 
-        cmd = f"docker run --name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.success_handler"
+        cmd = f"--name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.success_handler"
 
-        r = self.create_container_and_invoke_function(cmd, port)
+        with self.create_container(cmd, image):
+            r = self.invoke_function(port)
         
-        self.assertEqual(b'"My lambda ran succesfully"', r.content)
+            self.assertEqual(b'"My lambda ran succesfully"', r.content)
 
-    @parameterized.expand([("x86_64", "8010"), ("arm64", "9010"), ("", "9060")])
-    def test_function_name_is_overriden(self, arch, port):
+
+    def test_function_name_is_overriden(self, arch=ARCH, port=PORT):
         image, rie, image_name = self.tagged_name("assert-overwritten", arch)
 
-        cmd = f"docker run --name {image} -d --env AWS_LAMBDA_FUNCTION_NAME=MyCoolName -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.assert_env_var_is_overwritten"
+        cmd = f"--name {image} -d --env AWS_LAMBDA_FUNCTION_NAME=MyCoolName -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8080 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.assert_env_var_is_overwritten"
 
-        r = self.create_container_and_invoke_function(cmd, port)
+        with self.create_container(cmd, image):
+            r = self.invoke_function(port)
         
-        self.assertEqual(b'"My lambda ran succesfully"', r.content)
+            self.assertEqual(b'"My lambda ran succesfully"', r.content)
+
 
-    @parameterized.expand([("x86_64", "8011"), ("arm64", "9011"), ("", "9061")])
-    def test_port_override(self, arch, port):
+    def test_port_override(self, arch=ARCH, port=PORT):
         image, rie, image_name = self.tagged_name("port_override", arch)
 
         # Use port 8081 inside the container instead of 8080
-        cmd = f"docker run --name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8081 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.success_handler --runtime-interface-emulator-address 0.0.0.0:8081"
+        cmd = f"--name {image} -d -v {self.path_to_binary}:/local-lambda-runtime-server -p {port}:8081 --entrypoint /local-lambda-runtime-server/{rie} {image_name} {DEFAULT_1P_ENTRYPOINT} main.success_handler --runtime-interface-emulator-address 0.0.0.0:8081"
 
-        r = self.create_container_and_invoke_function(cmd, port)
+        with self.create_container(cmd, image):
+            r = self.invoke_function(port)
         
-        self.assertEqual(b'"My lambda ran succesfully"', r.content)
+            self.assertEqual(b'"My lambda ran succesfully"', r.content)
+
 
 
 if __name__ == "__main__":
diff --git a/test/integration/testdata/Dockerfile-allinone b/test/integration/testdata/Dockerfile-allinone
