schema: save original node req for reload.

The original node req shouldn't be changed in following logic to keep it
original state.

lisa/schema.py

@@ -1111,9 +1111,21 @@ class Environment:
     )
     nodes_requirement: Optional[List[NodeSpace]] = None
 
+    _original_nodes_requirement: Optional[List[NodeSpace]] = None
+
     def __post_init__(self, *args: Any, **kwargs: Any) -> None:
+        self._original_nodes_requirement = self.nodes_requirement
+        self.reload_requirements()
+
+    def reload_requirements(self) -> None:
         results: List[Node] = []
 
+        self.nodes = []
+        self.nodes_requirement = None
+        if self._original_nodes_requirement:
+            self.nodes_requirement = []
+            self.nodes_requirement.extend(copy.copy(self._original_nodes_requirement))
+
         if self.nodes_raw:
             for node_raw in self.nodes_raw:
                 node_type = node_raw[constants.TYPE]
@@ -1127,7 +1139,6 @@ class Environment:
                     # load base schema for future parsing
                     node: Node = load_by_type(Node, node_raw)
                     results.append(node)
-            self.nodes_raw = None
 
         self.nodes = results
 

selftests/azure/test_prepare.py

@@ -250,6 +250,7 @@ class AzurePrepareTestCase(TestCase):
         env.runbook.nodes_requirement.append(
             schema.NodeSpace(memory_mb=search_space.IntRange(min=143360))
         )
+        self._set_nodes_raw(env)
         self.verify_prepared_nodes(
             expected_result=True,
             expected_locations=["eastus2", "eastus2"],
@@ -265,6 +266,7 @@ class AzurePrepareTestCase(TestCase):
         env.runbook.nodes_requirement.append(
             schema.NodeSpace(core_count=8, memory_mb=16384)
         )
+        self._set_nodes_raw(env)
         self.verify_prepared_nodes(
             expected_result=True,
             expected_locations=["eastus2", "eastus2"],
@@ -283,6 +285,7 @@ class AzurePrepareTestCase(TestCase):
                 network_interface=schema.NetworkInterfaceOptionSettings(nic_count=3)
             )
         )
+        self._set_nodes_raw(env)
         self.verify_prepared_nodes(
             expected_result=True,
             expected_locations=["eastus2", "eastus2"],
@@ -341,6 +344,7 @@ class AzurePrepareTestCase(TestCase):
         environment = Environment(
             is_predefined=True, warn_as_error=False, id_=0, runbook=runbook
         )
+        self._set_nodes_raw(environment)
 
         return environment
 
@@ -357,6 +361,7 @@ class AzurePrepareTestCase(TestCase):
         ].get_extended_runbook(common.AzureNodeSchema, AZURE)
         node_runbook.location = location
         node_runbook.vm_size = vm_size
+        self._set_nodes_raw(environment)
 
     def verify_prepared_nodes(
         self,
@@ -405,3 +410,20 @@ class AzurePrepareTestCase(TestCase):
                 self.assertLessEqual(0, node_cap.gpu_count)
 
         self.assertEqual(expected_cost, environment.cost)
+
+    def _set_nodes_raw(self, environment: Environment) -> None:
+        if not environment.runbook.nodes_requirement:
+            return
+        environment.runbook.nodes_raw = []
+        for requirement in environment.runbook.nodes_requirement:
+            extended_runbook = requirement.get_extended_runbook(
+                common.AzureNodeSchema, AZURE
+            )
+            node_raw = requirement.to_dict()  # type: ignore
+            extended_raw = extended_runbook.to_dict()  # type: ignore
+            for field in ["marketplace", "shared_gallery"]:
+                if field in extended_raw and not extended_raw[field]:
+                    del extended_raw[field]
+            node_raw["azure"] = extended_raw
+
+            environment.runbook.nodes_raw.append(node_raw)

selftests/test_platform.py

@@ -153,9 +153,10 @@ class PlatformTestCase(TestCase):
         with self.assertRaises(LisaException) as cm:
             [platform.prepare_environment(env) for env in envs.values()]
         self.assertEqual(
-            "no capability found for environment: "
-            "Environment(name='customized_0', topology='subnet', "
-            "nodes_raw=None, nodes_requirement=None)",
+            "no capability found for environment: Environment("
+            "name='customized_0', topology='subnet', nodes_raw=[{'type': 'local', "
+            "'capability': {'core_count': {'min': 4}}}], nodes_requirement=None, "
+            "_original_nodes_requirement=None)",
             str(cm.exception),
         )