diff --git a/tests_e2e/tests/lib/network_security_rule.py b/tests_e2e/tests/lib/network_security_rule.py index d2f67d19cb..6fe8ee296e 100644 --- a/tests_e2e/tests/lib/network_security_rule.py +++ b/tests_e2e/tests/lib/network_security_rule.py @@ -55,7 +55,7 @@ def add_security_rule(self, security_rule: Dict[str, Any]) -> None: self._get_network_security_group()["properties"]["securityRules"].append(security_rule) def _get_network_security_group(self) -> Dict[str, Any]: - resources: Dict[str, Dict[str, Any]] = self._template["resources"] + resources: Any = self._template["resources"] # # If the NSG already exists, just return it # @@ -76,8 +76,14 @@ def _get_network_security_group(self) -> Dict[str, Any]: "securityRules": [] }} }}""") - nsg_reference = "network_security_groups" - resources[nsg_reference] = network_security_group + + # resources is a dictionary in LISA's ARM template, but a list in the template for scale sets + if isinstance(resources, dict): + nsg_reference = "network_security_groups" + resources[nsg_reference] = network_security_group + else: + nsg_reference = f"[resourceId('Microsoft.Network/networkSecurityGroups', '{self._NETWORK_SECURITY_GROUP}')]" + resources.append(network_security_group) # # Add a dependency on the NSG to the virtual network diff --git a/tests_e2e/tests/lib/update_arm_template.py b/tests_e2e/tests/lib/update_arm_template.py index 1cd4e1d72f..ef3dfd1d92 100644 --- a/tests_e2e/tests/lib/update_arm_template.py +++ b/tests_e2e/tests/lib/update_arm_template.py @@ -32,25 +32,29 @@ def update(self, template: Dict[str, Any], is_lisa_template: bool) -> None: """ @staticmethod - def get_resource(resources: Dict[str, Dict[str, Any]], type_name: str) -> Any: + def get_resource(resources: Any, type_name: str) -> Any: """ - Returns the first resource of the specified type in the given 'resources' list. + Returns the first resource of the specified type in the given 'resources' list/dict. Raises KeyError if no resource of the specified type is found. """ - for item in resources.values(): + if isinstance(resources, dict): + resources = resources.values() + for item in resources: if item["type"] == type_name: return item raise KeyError(f"Cannot find a resource of type {type_name} in the ARM template") @staticmethod - def get_resource_by_name(resources: Dict[str, Dict[str, Any]], resource_name: str, type_name: str) -> Any: + def get_resource_by_name(resources: Any, resource_name: str, type_name: str) -> Any: """ - Returns the first resource of the specified type and name in the given 'resources' list. + Returns the first resource of the specified type and name in the given 'resources' list/dict. Raises KeyError if no resource of the specified type and name is found. """ - for item in resources.values(): + if isinstance(resources, dict): + resources = resources.values() + for item in resources: if item["type"] == type_name and item["name"] == resource_name: return item raise KeyError(f"Cannot find a resource {resource_name} of type {type_name} in the ARM template")