Add hook azure_update_vm_capabilities

Add hook `azure_update_vm_capabilities` to allow editing of Azure VM SKU capabilities before being saved to cache.

This will allow capability information to be inserted by extensions for internal SKUs or the processing of additional SKU information such as disk and network performance. Alternatively, SKUs can be removed from the results to restrict the SKUs that LISA will deploy.

Author: kamalca

lisa/sut_orchestrator/azure/hooks.py

@@ -1,8 +1,9 @@
 import re
 from functools import partial
-from typing import Any, List, Pattern, Tuple
+from typing import Any, Dict, List, Pattern, Tuple
 
 from lisa.environment import Environment
+from lisa.sut_orchestrator.azure.common import AzureCapability
 from lisa.util import (
     ResourceAwaitableException,
     SkippedException,
@@ -34,6 +35,19 @@ def azure_update_arm_template(
         """
         ...
 
+    @hookspec
+    def azure_update_vm_capabilities(
+        self, capabilities: Dict[str, AzureCapability]
+    ) -> None:
+        """
+        Implement it to update the vm capabilities.
+
+        Args:
+            capabilities: the dict object mapping VM SKU name to Azure capability,
+                which is compiled from the output of _resource_sku_to_capability()
+        """
+        ...
+
 
 class AzureHookSpecDefaultImpl:
     __error_maps: List[Tuple[str, Pattern[str], Any]] = [

lisa/sut_orchestrator/azure/platform_.py

@@ -1111,6 +1111,7 @@ def get_location_info(self, location: str, log: Logger) -> AzureLocation:
                     except Exception as e:
                         log.error(f"unknown sku: {sku_obj}")
                         raise e
+            plugin_manager.hook.azure_update_vm_capabilities(capabilities=all_skus)
             location_data = AzureLocation(location=location, capabilities=all_skus)
             log.debug(f"{location}: saving to disk")
             with open(cached_file_name, "w") as f: