extract vm set up into a method

microsoft/testsuites/vm_extensions/linux_patch_extension.py

@@ -1,6 +1,6 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT license.
+# Copyright (c) Microsoft Corporation. Licensed under the MIT license.
 
+from typing import Any
 
 from assertpy.assertpy import assert_that
 
@@ -22,15 +22,30 @@
 from lisa.sut_orchestrator.azure.platform_ import AzurePlatform
 
 
+def _set_up_vm(node: Node, environment: Environment) -> Any:
+    assert environment.platform, "platform shouldn't be None."
+    platform: AzurePlatform = environment.platform  # type: ignore
+    assert isinstance(
+        platform, AzurePlatform
+    ), "platform should be AzurePlatform instance"
+    assert isinstance(
+        platform, AzurePlatform
+    ), "platform should be AzurePlatform instance"
+    compute_client = get_compute_client(platform)
+    node_context = get_node_context(node)
+    resource_group_name = node_context.resource_group_name
+    vm_name = node_context.vm_name
+
+    return compute_client, resource_group_name, vm_name
+
+
 def _verify_vm_agent_running(node: Node, log: Logger) -> None:
     service = node.tools[Service]
     is_vm_agent_running = service.is_service_running(
         "walinuxagent.service"
     ) or service.is_service_running("waagent.service")
 
-    log.debug(
-        f"verify walinuxagent or waagent service is running: {is_vm_agent_running}"
-    )
+    log.debug(f"verify walinuxagent or waagent running:{is_vm_agent_running}")
 
     assert_that(is_vm_agent_running).described_as(
         "Expected walinuxagent or waagent service is running"
@@ -46,32 +61,25 @@ def _verify_vm_agent_running(node: Node, log: Logger) -> None:
 class LinuxPatchExtensionBVT(TestSuite):
     @TestCaseMetadata(
         description="""
-        Verify walinuxagent or waagent service is running on virtual machine.
-        Perform assess patches to trigger Microsoft.CPlat.Core.LinuxPatchExtension creation in a virtual machine.
-        Verify status file response for validity.
+        Verify walinuxagent or waagent service is running on vm. Perform assess
+        patches to trigger Microsoft.CPlat.Core.LinuxPatchExtension creation in
+        vm. Verify status file response for validity.
         """,
         priority=1,
     )
     def verify_vm_assess_patches(
         self, node: Node, environment: Environment, log: Logger
     ) -> None:
-        assert environment.platform, "platform shouldn't be None."
-        platform: AzurePlatform = environment.platform  # type: ignore
-        assert isinstance(
-            platform, AzurePlatform
-        ), "platform should be AzurePlatform instance"
-        compute_client = get_compute_client(platform)
-        node_context = get_node_context(node)
-        resource_group_name = node_context.resource_group_name
-        vm_name = node_context.vm_name
-
-        # verify vm agent is running
+        compute_client, resource_group_name, vm_name = _set_up_vm(node, environment)
+        # verify vm agent service is running, lpe is a dependent of vm agent
+        # service
         _verify_vm_agent_running(node, log)
 
         operation = compute_client.virtual_machines.begin_assess_patches(
             resource_group_name=resource_group_name, vm_name=vm_name
         )
-        # set wait operation timeout 10 min, status file should be generated before timeout
+        # set wait operation timeout 10 min, status file should be generated
+        # before timeout
         assess_result = wait_operation(operation, 600)
 
         assert assess_result, "assess_result shouldn't be None"
@@ -85,24 +93,16 @@ def verify_vm_assess_patches(
 
     @TestCaseMetadata(
         description="""
-        Verify walinuxagent or waagent service is running on virtual machine.
-        Perform install patches to trigger Microsoft.CPlat.Core.LinuxPatchExtension creation in a virtual machine.
+        Verify walinuxagent or waagent service is running on vm. Perform install
+        patches to trigger Microsoft.CPlat.Core.LinuxPatchExtension creation in vm.
         Verify status file response for validity.
         """,
         priority=2,
     )
     def verify_vm_install_patches(
         self, node: Node, environment: Environment, log: Logger
     ) -> None:
-        assert environment.platform, "platform shouldn't be None."
-        platform: AzurePlatform = environment.platform  # type: ignore
-        assert isinstance(
-            platform, AzurePlatform
-        ), "platform should be AzurePlatform instance"
-        compute_client = get_compute_client(platform)
-        node_context = get_node_context(node)
-        resource_group_name = node_context.resource_group_name
-        vm_name = node_context.vm_name
+        compute_client, resource_group_name, vm_name = _set_up_vm(node, environment)
         install_patches_input = {
             "maximumDuration": "PT3H30M",
             "rebootSetting": "IfRequired",
@@ -112,22 +112,24 @@ def verify_vm_install_patches(
             },
         }
 
-        # verify vm agent is running
+        # verify vm agent service is running, lpe is a dependent of vm agent
+        # service
         _verify_vm_agent_running(node, log)
 
         operation = compute_client.virtual_machines.begin_install_patches(
             resource_group_name=resource_group_name,
             vm_name=vm_name,
             install_patches_input=install_patches_input,
         )
-        # set wait operation timeout 10 min, status file should be generated before timeout
+        # set wait operation timeout 10 min, status file should be generated
+        # before timeout
         install_result = wait_operation(operation, 600)
 
-        assert install_result, "assess_result shouldn't be None"
+        assert install_result, "install_result shouldn't be None"
         assert_that(install_result["status"]).described_as(
-            "Expected the assess patches to succeed"
+            "Expected the install patches to succeed"
         ).is_equal_to("Succeeded")
 
         assert_that(install_result["error"]["code"]).described_as(
-            "Expected no error in assess patches operation"
+            "Expected no error in install patches operation"
         ).is_equal_to("0")