Allow tests to run on random images

tests_e2e/orchestrator/lib/agent_test_loader.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 #
 import importlib.util
+import re
 # E0401: Unable to import 'yaml' (import-error)
 import yaml  # pylint: disable=E0401
 
@@ -118,13 +119,19 @@ def images(self) -> Dict[str, List[VmImageInfo]]:
         """
         return self.__images
 
+    # Matches a reference to a random subset of images within a set with an optional count: random(<image_set>, [<count>]), e.g. random(endorsed, 3), random(endorsed)
+    RANDOM_IMAGES_RE = re.compile(r"random\((?P<image_set>[^,]+)(\s*,\s*(?P<count>\d+))?\)")
+
     def _validate(self):
         """
         Performs some basic validations on the data loaded from the YAML description files
         """
         for suite in self.test_suites:
             # Validate that the images the suite must run on are in images.yml
             for image in suite.images:
+                match = AgentTestLoader.RANDOM_IMAGES_RE.match(image)
+                if match is not None:
+                    image = match.group('image_set')
                 if image not in self.images:
                     raise Exception(f"Invalid image reference in test suite {suite.name}: Can't find {image} in images.yml")
 

tests_e2e/orchestrator/lib/agent_test_suite_combinator.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 import logging
+import random
 import re
 import urllib.parse
 
@@ -17,7 +18,7 @@
 from lisa.combinator import Combinator  # pylint: disable=E0401
 from lisa.util import field_metadata  # pylint: disable=E0401
 
-from tests_e2e.orchestrator.lib.agent_test_loader import AgentTestLoader, VmImageInfo
+from tests_e2e.orchestrator.lib.agent_test_loader import AgentTestLoader, VmImageInfo, TestSuiteInfo
 
 
 @dataclass_json()
@@ -118,160 +119,215 @@ def _next(self) -> Optional[Dict[str, Any]]:
     def create_environment_for_existing_vm(self) -> List[Dict[str, Any]]:
         loader = AgentTestLoader(self.runbook.test_suites, self.runbook.cloud)
 
-        environment: List[Dict[str, Any]] = [
-            {
-                "c_env_name": self.runbook.vm_name,
-                "c_vm_name": self.runbook.vm_name,
-                "c_location": self.runbook.location,
-                "c_test_suites": loader.test_suites,
-            }
-        ]
+        environment: Dict[str, Any] = {
+            "c_env_name": self.runbook.vm_name,
+            "c_vm_name": self.runbook.vm_name,
+            "c_location": self.runbook.location,
+            "c_test_suites": loader.test_suites,
+        }
 
         log: logging.Logger = logging.getLogger("lisa")
-        log.info("******** Environment for existing VMs *****")
-        log.info(
-            "{  c_env_name: '%s', c_vm_name: '%s', c_location: '%s', c_test_suites: '%s' }",
-            environment[0]['c_env_name'], environment[0]['c_vm_name'], environment[0]['c_location'], [s.name for s in environment[0]['c_test_suites']])
-        log.info("***************************")
+        log.info("******** Waagent: Settings for existing VM *****")
+        log.info("")
+        log.info("Settings for %s:\n%s\n", environment['c_env_name'], self._get_env_settings(environment))
+        log.info("")
 
-        return environment
+        return [environment]
 
     def create_environment_list(self) -> List[Dict[str, Any]]:
+        """
+        Examines the test_suites specified in the runbook and returns a list of the environments (i.e. test VMs) that need to be
+        created in order to execute these suites.
+
+        Note that if the runbook provides an 'image', 'location', or 'vm_size', those values override any values provided in the
+        configuration of the test suites.
+        """
+        environments: List[Dict[str, Any]] = []
+        shared_environments: Dict[str, Dict[str, Any]] = {}  # environments shared by multiple test suites
+
         loader = AgentTestLoader(self.runbook.test_suites, self.runbook.cloud)
 
-        #
-        # If the runbook provides any of 'image', 'location', or 'vm_size', those values
-        # override any configuration values on the test suite.
-        #
-        # Check 'images' first and add them to 'runbook_images', if any
-        #
-        if self.runbook.image == "":
-            runbook_images = []
-        else:
-            runbook_images = loader.images.get(self.runbook.image)
-            if runbook_images is None:
-                if not self._is_urn(self.runbook.image) and not self._is_vhd(self.runbook.image):
-                    raise Exception(f"The 'image' parameter must be an image or image set name, a urn, or a vhd: {self.runbook.image}")
-                i = VmImageInfo()
-                i.urn = self.runbook.image  # Note that this could be a URN or the URI for a VHD
-                i.locations = []
-                i.vm_sizes = []
-                runbook_images = [i]
-
-        #
-        # Now walk through all the test_suites and create a list of the environments (test VMs) that need to be created.
-        #
-        environment_list: List[Dict[str, Any]] = []
-        shared_environments: Dict[str, Dict[str, Any]] = {}
+        runbook_images = self._get_runbook_images(loader)
 
         for suite_info in loader.test_suites:
             if len(runbook_images) > 0:
-                images_info = runbook_images
+                images_info: List[VmImageInfo] = runbook_images
             else:
-                # The test suite may be referencing multiple image sets, and sets can intersect, so we need to ensure
-                # we eliminate any duplicates.
-                unique_images: Dict[str, str] = {}
-                for image in suite_info.images:
-                    for i in loader.images[image]:
-                        unique_images[i] = i
-                images_info = unique_images.values()
+                images_info: List[VmImageInfo] = self._get_test_suite_images(suite_info, loader)
 
             for image in images_info:
-                # The URN can actually point to a VHD if the runbook provided a VHD in the 'images' parameter
+                # 'image.urn' can actually be the URL to a VHD if the runbook provided it in the 'image' parameter
                 if self._is_vhd(image.urn):
-                    marketplace_image = ""
-                    vhd = image.urn
-                    name = "vhd"
+                    c_marketplace_image = ""
+                    c_vhd = image.urn
+                    image_name = "vhd"
                 else:
-                    marketplace_image = image.urn
-                    vhd = ""
-                    match = AgentTestSuitesCombinator._URN.match(image.urn)
-                    if match is None:
-                        raise Exception(f"Invalid URN: {image.urn}")
-                    name = f"{match.group('offer')}-{match.group('sku')}"
-
-                location: str = None
-                # If the runbook specified a location, use it.
-                if self.runbook.location != "":
-                    location = self.runbook.location
-                #  Then try the suite location, if any.
-                elif suite_info.location != '':
-                    location = suite_info.location
-                # If the image has a location restriction, use any location where it is available.
-                # However, if it is not available on any location, skip the image.
-                elif image.locations:
-                    image_locations = image.locations.get(self.runbook.cloud)
-                    if image_locations is not None:
-                        if len(image_locations) == 0:
-                            continue
-                        location = image_locations[0]
-                # If no location has been selected, use the default.
-                if location is None:
-                    location = AgentTestSuitesCombinator._DEFAULT_LOCATIONS[self.runbook.cloud]
-
-                # If the runbook specified a VM size, use it. Else if the image specifies a list of VM sizes, use any of them. Otherwise,
-                # set the size to empty and let LISA choose it.
-                if self.runbook.vm_size != '':
-                    vm_size = self.runbook.vm_size
-                elif len(image.vm_sizes) > 0:
-                    vm_size = image.vm_sizes[0]
-                else:
-                    vm_size = ""
+                    c_marketplace_image = image.urn
+                    c_vhd = ""
+                    image_name = self._get_image_name(image.urn)
+
+                c_location: str = self._get_location(suite_info, image)
+                if c_location is None:
+                    continue
+
+                c_vm_size = self._get_vm_size(image)
 
                 # Note: Disabling "W0640: Cell variable 'foo' defined in loop (cell-var-from-loop)". This is a false positive, the closure is OK
                 # to use, since create_environment() is called within the same iteration of the loop.
                 # pylint: disable=W0640
-                def create_environment(env_name: str) -> Dict[str, Any]:
-                    tags = {}
+                def create_environment(c_env_name: str) -> Dict[str, Any]:
+                    c_vm_tags = {}
                     if suite_info.template != '':
-                        tags["templates"] = suite_info.template
+                        c_vm_tags["templates"] = suite_info.template
                     return {
-                        "c_marketplace_image": marketplace_image,
+                        "c_marketplace_image": c_marketplace_image,
                         "c_cloud": self.runbook.cloud,
-                        "c_location": location,
-                        "c_vm_size": vm_size,
-                        "c_vhd": vhd,
+                        "c_location": c_location,
+                        "c_vm_size": c_vm_size,
+                        "c_vhd": c_vhd,
                         "c_test_suites": [suite_info],
-                        "c_env_name": env_name,
+                        "c_env_name": c_env_name,
                         "c_marketplace_image_information_location": self._MARKETPLACE_IMAGE_INFORMATION_LOCATIONS[self.runbook.cloud],
                         "c_shared_resource_group_location": self._SHARED_RESOURCE_GROUP_LOCATIONS[self.runbook.cloud],
-                        "c_vm_tags": tags
+                        "c_vm_tags": c_vm_tags
                     }
                 # pylint: enable=W0640
 
                 if suite_info.owns_vm:
                     # create an environment for exclusive use by this suite
-                    environment_list.append(create_environment(f"{name}-{suite_info.name}"))
+                    environments.append(create_environment(f"{image_name}-{suite_info.name}"))
                 else:
                     # add this suite to the shared environments
-                    key: str = f"{name}-{location}"
-                    environment = shared_environments.get(key)
-                    if environment is not None:
-                        environment["c_test_suites"].append(suite_info)
+                    key: str = f"{image_name}-{c_location}"
+                    env = shared_environments.get(key)
+                    if env is not None:
+                        env["c_test_suites"].append(suite_info)
                         if suite_info.template != '':
-                            vm_tags = environment["c_vm_tags"]
+                            vm_tags = env["c_vm_tags"]
                             if "templates" in vm_tags:
                                 vm_tags["templates"] += ", " + suite_info.template
                             else:
                                 vm_tags["templates"] = suite_info.template
                     else:
                         shared_environments[key] = create_environment(key)
 
-        environment_list.extend(shared_environments.values())
+        environments.extend(shared_environments.values())
 
-        if len(environment_list) == 0:
+        if len(environments) == 0:
             raise Exception("No VM images were found to execute the test suites.")
 
         log: logging.Logger = logging.getLogger("lisa")
         log.info("")
         log.info("******** Waagent: Test Environments *****")
         log.info("")
-        for environment in environment_list:
-            test_suites = [s.name for s in environment['c_test_suites']]
-            log.info("Settings for %s:\n%s\n", environment['c_env_name'], '\n'.join([f"\t{name}: {value if name != 'c_test_suites' else test_suites}" for name, value in environment.items()]))
+        for env in environments:
+            log.info("Settings for %s:\n%s\n", env['c_env_name'], self._get_env_settings(env))
         log.info("")
 
-        return environment_list
+        return environments
+
+    def _get_runbook_images(self, loader: AgentTestLoader) -> List[VmImageInfo]:
+        """
+        Returns the images specified in the runbook, or an empty list if none are specified.
+        """
+        if self.runbook.image == "":
+            return []
+
+        images = loader.images.get(self.runbook.image)
+        if images is not None:
+            return images
+
+        # If it is not image or image set, it must be a URN or VHD
+        if not self._is_urn(self.runbook.image) and not self._is_vhd(self.runbook.image):
+            raise Exception(f"The 'image' parameter must be an image, an image set name, a urn, or a vhd: {self.runbook.image}")
+
+        i = VmImageInfo()
+        i.urn = self.runbook.image  # Note that this could be a URN or the URI for a VHD
+        i.locations = []
+        i.vm_sizes = []
+
+        return [i]
+
+    @staticmethod
+    def _get_test_suite_images(suite: TestSuiteInfo, loader: AgentTestLoader) -> List[VmImageInfo]:
+        """
+        Returns the images used by a test suite.
+
+        A test suite may be reference multiple image sets and sets can intersect; this method eliminates any duplicates.
+        """
+        unique: Dict[str, VmImageInfo] = {}
+        for image in suite.images:
+            match = AgentTestLoader.RANDOM_IMAGES_RE.match(image)
+            if match is None:
+                image_list = loader.images[image]
+            else:
+                count = match.group('count')
+                if count is None:
+                    count = 1
+                matching_images = loader.images[match.group('image_set')].copy()
+                random.shuffle(matching_images)
+                image_list = matching_images[0:count]
+            for i in image_list:
+                unique[i.urn] = i
+        return [v for k, v in unique.items()]
+
+    def _get_location(self, suite_info: TestSuiteInfo, image: VmImageInfo) -> str:
+        """
+        Returns the location on which the test VM for the given test suite and image should be created.
+
+        If the image is not available on any location, returns None, to indicate that the test suite should be skipped.
+        """
+        # If the runbook specified a location, use it.
+        if self.runbook.location != "":
+            return self.runbook.location
+
+        #  Then try the suite location, if any.
+        if suite_info.location != '':
+            return suite_info.location
+
+        # If the image has a location restriction, use any location where it is available.
+        # However, if it is not available on any location, skip the image (return None)
+        if image.locations:
+            image_locations = image.locations.get(self.runbook.cloud)
+            if image_locations is not None:
+                if len(image_locations) == 0:
+                    return None
+                return image_locations[0]
+
+        # Else use the default.
+        return AgentTestSuitesCombinator._DEFAULT_LOCATIONS[self.runbook.cloud]
+
+    def _get_vm_size(self, image: VmImageInfo) -> str:
+        """
+        Returns the VM size that should be used to create the test VM for the given image.
+
+        If the size is set to an empty string, LISA will choose an appropriate size
+        """
+        # If the runbook specified a VM size, use it.
+        if self.runbook.vm_size != '':
+            return self.runbook.vm_size
+
+        #  If the image specifies a list of VM sizes, use any of them.
+        if len(image.vm_sizes) > 0:
+            return image.vm_sizes[0]
+
+        # Otherwise, set the size to empty and LISA will select an appropriate size.
+        return ""
+
+    @staticmethod
+    def _get_image_name(urn: str) -> str:
+        """
+        Creates an image name ("offer-sku") given its URN
+        """
+        match = AgentTestSuitesCombinator._URN.match(urn)
+        if match is None:
+            raise Exception(f"Invalid URN: {urn}")
+        return f"{match.group('offer')}-{match.group('sku')}"
+
+    @staticmethod
+    def _get_env_settings(environment: Dict[str, Any]):
+        suite_names = [s.name for s in environment['c_test_suites']]
+        return '\n'.join([f"\t{name}: {value if name != 'c_test_suites' else suite_names}" for name, value in environment.items()])
 
     _URN = re.compile(r"(?P<publisher>[^\s:]+)[\s:](?P<offer>[^\s:]+)[\s:](?P<sku>[^\s:]+)[\s:](?P<version>[^\s:]+)")
 

tests_e2e/orchestrator/runbook.yml

@@ -49,7 +49,7 @@ variable:
   #
   # The test suites to execute
   - name: test_suites
-    value: "agent_bvt"
+    value: "agent_bvt, no_outbound_connections"
   - name: cloud
     value: "AzureCloud"
   - name: image

tests_e2e/test_suites/no_outbound_connections.yml

@@ -15,7 +15,6 @@ tests:
   - "bvts/run_command.py"
   - "bvts/vm_access.py"
   - "no_outbound_connections/check_fallback_to_hgap.py"
-images:
-  - "ubuntu_2004"
+images: "random(endorsed)"
 template: "no_outbound_connections/nsg_template.py"
 owns_vm: true