From a409eba8160f8f8ac00c7de2a853caea16fc3239 Mon Sep 17 00:00:00 2001
From: SrikanthMyakam <m.srikanth.m@gmail.com>
Date: Sat, 25 Jan 2025 13:33:45 +0530
Subject: [PATCH] Improvements to "HyperVPreparationTransformer" (#3613)

* Improvements to "HyperVPreparationTransformer"

Changed the transformer to use recently added methods to HyperV tool instead of running commands.

* Update hyperv_preparation.py

* Update hyperv.py
---
 lisa/tools/hyperv.py                    | 11 ++++++---
 lisa/transformers/hyperv_preparation.py | 33 ++++++++++---------------
 2 files changed, 20 insertions(+), 24 deletions(-)

diff --git a/lisa/tools/hyperv.py b/lisa/tools/hyperv.py
index e398939b4e..d328aa72f5 100644
--- a/lisa/tools/hyperv.py
+++ b/lisa/tools/hyperv.py
@@ -232,9 +232,12 @@ def get_default_switch(self) -> VMSwitch:
                 raise LisaException("Could not find any Internal or External switch")
         return self._default_switch
 
-    def exists_switch(self, name: str) -> bool:
+    def exists_switch(self, name: str, switch_type: str = "") -> bool:
+        cmd = f"Get-VMSwitch -Name {name}"
+        if switch_type != "":
+            cmd += f"  -SwitchType '{switch_type}'"
         output = self.node.tools[PowerShell].run_cmdlet(
-            f"Get-VMSwitch -Name {name}",
+            cmdlet=cmd,
             fail_on_error=False,
             force_run=True,
         )
@@ -248,8 +251,8 @@ def delete_switch(self, name: str) -> None:
             )
 
     def create_switch(self, name: str, switch_type: str = "Internal") -> None:
-        # remove switch if it exists
-        self.delete_switch(name)
+        if self.exists_switch(name, switch_type):
+            return
 
         # create a new switch
         self.node.tools[PowerShell].run_cmdlet(
diff --git a/lisa/transformers/hyperv_preparation.py b/lisa/transformers/hyperv_preparation.py
index e2ec20b7ca..f2296728c0 100644
--- a/lisa/transformers/hyperv_preparation.py
+++ b/lisa/transformers/hyperv_preparation.py
@@ -1,7 +1,7 @@
 from typing import Any, Dict, List, Type
 
 from lisa import schema
-from lisa.tools import PowerShell
+from lisa.tools import HyperV
 from lisa.transformers.deployment_transformer import (
     DeploymentTransformer,
     DeploymentTransformerSchema,
@@ -28,23 +28,16 @@ def _output_names(self) -> List[str]:
     def _internal_run(self) -> Dict[str, Any]:
         runbook: DeploymentTransformerSchema = self.runbook
         assert isinstance(runbook, DeploymentTransformerSchema)
-        node = self._node
-        powershell = node.tools[PowerShell]
-        powershell.run_cmdlet(
-            "Install-WindowsFeature -Name DHCP,Hyper-V  -IncludeManagementTools",
-            force_run=True,
-        )
-        node.reboot()
-        powershell.run_cmdlet(
-            "New-VMSwitch -Name 'InternalNAT' -SwitchType Internal",
-            force_run=True,
-        )
-        powershell.run_cmdlet(
-            "New-NetNat -Name 'InternalNAT' -InternalIPInterfaceAddressPrefix '192.168.0.0/24'",  # noqa: E501
-            force_run=True,
-        )
-        powershell.run_cmdlet(
-            'New-NetIPAddress -IPAddress 192.168.0.1 -InterfaceIndex (Get-NetAdapter | Where-Object { $_.Name -like "*InternalNAT)" } | Select-Object -ExpandProperty ifIndex) -PrefixLength 24',  # noqa: E501
-            force_run=True,
-        )
+        switch_name = "InternalNAT"
+
+        # Enable Hyper-V
+        hv = self._node.tools[HyperV]
+
+        # Create an internal switch.
+        hv.create_switch(name=switch_name)
+
+        hv.setup_nat_networking(switch_name=switch_name, nat_name=switch_name)
+
+        # Configure Internal DHCP
+        hv.enable_internal_dhcp()
         return {}