From d77d7739e293b5ad4c7c43b448636a6a42d08cb0 Mon Sep 17 00:00:00 2001
From: "Chi Song (from Dev Box)" <chisong@microsoft.com>
Date: Tue, 8 Oct 2024 09:34:09 -0700
Subject: [PATCH] Hibernation: Change is_enabled to the search space
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Hibernation must be configured during deployment. If it’s not set up and the hibernation feature is available, hibernation test cases will fail. This update changes `is_enabled` to a search space. If hibernation is not configured, the value will default to `False`. For test cases requiring hibernation to be enabled, a new environment must be set up with the feature properly configured.
---
 lisa/features/hibernation.py            | 47 +++++++++++++++++++++----
 lisa/sut_orchestrator/azure/features.py |  3 +-
 2 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/lisa/features/hibernation.py b/lisa/features/hibernation.py
index b1b88676a..c4275aac6 100644
--- a/lisa/features/hibernation.py
+++ b/lisa/features/hibernation.py
@@ -1,22 +1,32 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from functools import partial
-from typing import Any, Type, cast
+from typing import Any, List, Type, Union, cast
 
 from dataclasses_json import dataclass_json
 
-from lisa import schema
+from lisa import schema, search_space
 from lisa.feature import Feature
+from lisa.util import field_metadata
 
 FEATURE_NAME_HIBERNATION = "Hibernation"
 
+false_priority: List[bool] = [False, True]
+
 
 @dataclass_json()
 @dataclass()
 class HibernationSettings(schema.FeatureSettings):
     type: str = FEATURE_NAME_HIBERNATION
-    is_enabled: bool = False
+    is_enabled: Union[search_space.SetSpace[bool], bool] = field(
+        default_factory=partial(
+            search_space.SetSpace[bool], is_allow_set=True, items=[True, False]
+        ),
+        metadata=field_metadata(
+            decoder=partial(search_space.decode_set_space_by_type, base_type=bool)
+        ),
+    )
 
     def __hash__(self) -> int:
         return hash(self._get_key())
@@ -24,8 +34,33 @@ class HibernationSettings(schema.FeatureSettings):
     def _get_key(self) -> str:
         return f"{self.type}/{self.is_enabled}"
 
-    def _generate_min_capability(self, capability: Any) -> Any:
-        return self
+    def _call_requirement_method(
+        self, method: search_space.RequirementMethod, capability: Any
+    ) -> Any:
+        assert isinstance(
+            capability, HibernationSettings
+        ), f"actual: {type(capability)}"
+
+        value = HibernationSettings()
+        value.is_enabled = getattr(
+            search_space, f"{method.value}_setspace_by_priority"
+        )(
+            self.is_enabled,
+            capability.is_enabled,
+            false_priority,
+        )
+        return value
+
+    def check(self, capability: Any) -> search_space.ResultReason:
+        assert isinstance(
+            capability, HibernationSettings
+        ), f"actual: {type(capability)}"
+        result = super().check(capability)
+        result.merge(
+            search_space.check_setspace(self.is_enabled, capability.is_enabled),
+            "is_enabled",
+        )
+        return result
 
 
 class Hibernation(Feature):
diff --git a/lisa/sut_orchestrator/azure/features.py b/lisa/sut_orchestrator/azure/features.py
index f6043a5a3..ec72e1283 100644
--- a/lisa/sut_orchestrator/azure/features.py
+++ b/lisa/sut_orchestrator/azure/features.py
@@ -42,6 +42,7 @@ from lisa.environment import Environment
 from lisa.feature import Feature
 from lisa.features.availability import AvailabilityType
 from lisa.features.gpu import ComputeSDK
+from lisa.features.hibernation import HibernationSettings
 from lisa.features.resize import ResizeAction
 from lisa.features.security_profile import (
     FEATURE_NAME_SECURITY_PROFILE,
@@ -2195,7 +2196,7 @@ class Hibernation(AzureFeatureMixin, features.Hibernation):
             ]
             or raw_capabilities.get("HibernationSupported", None) == "True"
         ):
-            return schema.FeatureSettings.create(cls.name())
+            return HibernationSettings()
 
         return None