making some tests run quicker

src/confcom/azext_confcom/tests/latest/test_confcom_scenario.py

@@ -28,7 +28,7 @@ class MountEnforcement(unittest.TestCase):
         "version": "1.0",
         "containers": [
             {
-                "containerImage": "rust:1.52.1",
+                "containerImage": "alpine:3.16",
                 "environmentVariables": [
                     {
                         "name": "PATH",
@@ -51,7 +51,7 @@ class MountEnforcement(unittest.TestCase):
                 ]
             },
             {
-                "containerImage": "python:3.6.14-slim-buster",
+                "containerImage": "nginx:1.24",
                 "environmentVariables": [],
                 "command": ["echo", "hello"],
                 "workingDir": "/customized/absolute/path",
@@ -76,7 +76,7 @@ def test_user_container_customized_mounts(self):
             (
                 img
                 for img in self.aci_policy.get_images()
-                if isinstance(img, UserContainerImage) and img.base == "rust"
+                if isinstance(img, UserContainerImage) and img.base == "alpine"
             ),
             None,
         )
@@ -115,7 +115,7 @@ def test_user_container_mount_injected_dns(self):
             (
                 img
                 for img in self.aci_policy.get_images()
-                if isinstance(img, UserContainerImage) and img.base == "python"
+                if isinstance(img, UserContainerImage) and img.base == "nginx"
             ),
             None,
         )

src/confcom/azext_confcom/tests/latest/test_confcom_tar.py

@@ -10,7 +10,8 @@
 import deepdiff
 import json
 import docker
-
+import shutil
+import threading
 from azext_confcom.security_policy import (
     OutputType,
     load_policy_from_arm_template_str,
@@ -20,7 +21,7 @@
 )
 import azext_confcom.config as config
 
-
+sem = threading.Semaphore(4)
 # @unittest.skip("not in use")
 @pytest.mark.run(order=11)
 class PolicyGeneratingArmParametersCleanRoomTarFile(unittest.TestCase):
@@ -33,6 +34,20 @@ def setUpClass(cls) -> None:
         cls.path = path
 
         cls.image_path = image_path
+        client = docker.from_env()
+        image = client.images.get("nginx:1.22")
+        f = open(image_path, "wb")
+        for chunk in image.save(named=True):
+            f.write(chunk)
+        f.close()
+        client.close()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        # delete the tar file once all the tests are done
+        if not sem.acquire(blocking=False) and os.path.isfile(cls.image_path):
+            os.remove(cls.image_path)
+
 
     def test_arm_template_with_parameter_file_clean_room_tar(self):
         custom_arm_json_default_value = """
@@ -169,27 +184,15 @@ def test_arm_template_with_parameter_file_clean_room_tar(self):
             custom_arm_json_default_value, ""
         )[0]
 
-        # save the tar file for the image in the testing directory
-        client = docker.from_env()
-        image = client.images.get("nginx:1.22")
-        tar_mapping_file = {"nginx:1.22": self.image_path}
-        # Note: Class setup and teardown shouldn't have side effects, and reading from the tar file fails when all the tests are running in parallel, so we want to save and delete this tar file as a part of the test. Not as a part of the testing class.
-        f = open(self.image_path, "wb")
-        for chunk in image.save(named=True):
-            f.write(chunk)
-        f.close()
-        client.close()
         tar_mapping_file = {"nginx:1.22": self.image_path}
         try:
             clean_room_image.populate_policy_content_for_all_images(
                 tar_mapping=tar_mapping_file
             )
-        except:
+            sem.release()
+        except Exception as e:
+            print(e)
             raise AccContainerError("Could not get image from tar file")
-        finally:
-            # delete the tar file
-            if os.path.isfile(self.image_path):
-                os.remove(self.image_path)
 
         regular_image_json = json.loads(
             regular_image.get_serialized_output(output_type=OutputType.RAW, rego_boilerplate=False)
@@ -384,32 +387,12 @@ def test_arm_template_mixed_mode_tar(self):
             custom_arm_json_default_value, ""
         )[0]
 
-        # save the tar file for the image in the testing directory
-        client = docker.from_env()
-        image = client.images.get("nginx:1.22")
-        image_path = self.image_path + "2"
-        # Note: Class setup and teardown shouldn't have side effects, and reading from the tar file fails when all the tests are running in parallel, so we want to save and delete this tar file as a part of the test. Not as a part of the testing class.
-        # make a temp directory for the tar file
-        temp_dir = tempfile.TemporaryDirectory()
+        image_path = self.image_path
 
-        image_path = os.path.join(
-            temp_dir.name, "nginx.tar"
-        )
-        f = open(image_path, "wb")
-        for chunk in image.save(named=True):
-            f.write(chunk)
-        f.close()
-        client.close()
-        tar_mapping_file = {"nginx:1.22": image_path}
-        try:
-            clean_room_image.populate_policy_content_for_all_images(
-                tar_mapping=image_path
+        clean_room_image.populate_policy_content_for_all_images(
+            tar_mapping=image_path
             )
-        finally:
-            temp_dir.cleanup()
-            # delete the tar file
-            if os.path.isfile(image_path):
-                os.remove(image_path)
+        sem.release()
 
         regular_image_json = json.loads(
             regular_image.get_serialized_output(output_type=OutputType.RAW, rego_boilerplate=False)
@@ -559,35 +542,16 @@ def test_arm_template_with_parameter_file_clean_room_tar_invalid(self):
         clean_room_image = load_policy_from_arm_template_str(
             custom_arm_json_default_value, ""
         )[0]
-        # save the tar file for the image in the testing directory
-        client = docker.from_env()
-        image = client.images.pull("nginx:1.23")
-        image = client.images.get("nginx:1.23")
-
-        # Note: Class setup and teardown shouldn't have side effects, and reading from the tar file fails when all the tests are running in parallel, so we want to save and delete this tar file as a part of the test. Not as a part of the testing class.
-        temp_dir = tempfile.TemporaryDirectory()
-
-        image_path = os.path.join(
-            temp_dir.name, "nginx.tar"
-        )
-        f = open(image_path, "wb")
-        for chunk in image.save(named=True):
-            f.write(chunk)
-        f.close()
-        client.close()
 
+        image_path = self.image_path
         try:
             clean_room_image.populate_policy_content_for_all_images(
                 tar_mapping=image_path
             )
+            sem.release()
             raise AccContainerError("getting image should fail")
         except:
             pass
-        finally:
-            # delete the tar file
-            temp_dir.cleanup()
-            if os.path.isfile(self.image_path):
-                os.remove(self.image_path)
 
     def test_clean_room_fake_tar_invalid(self):
         custom_arm_json_default_value = """
@@ -723,4 +687,5 @@ def test_clean_room_fake_tar_invalid(self):
             )
             raise AccContainerError("getting image should fail")
         except FileNotFoundError:
+            sem.release()
             pass

src/confcom/azext_confcom/tests/latest/test_confcom_template_util.py

@@ -489,16 +489,17 @@ def test_inject_policy_into_template(self):
             }
         }
         """
+        filename = "test_template.json"
         # write template to file for testing
-        with open("test_template.json", "w") as f:
+        with open(filename, "w") as f:
             f.write(template)
 
         with self.assertRaises(SystemExit) as exc_info:
-            acipolicygen_confcom(None, "test_template.json", None, None, None, None)
+            acipolicygen_confcom(None, filename, None, None, None, None)
 
         self.assertEqual(exc_info.exception.code, 0)
 
-        with open("test_template.json", "r") as f:
+        with open(filename, "r") as f:
             template_with_policy = load_json_from_str(f.read())
 
             # check if template contains confidential compute policy
@@ -528,4 +529,4 @@ def test_inject_policy_into_template(self):
                 > 0
             )
         # delete test file
-        os.remove("test_template.json")
+        os.remove(filename)