Implement Power Outage Rollback Feature (#927)

Signed-off-by: SK8-infi <shivansh.katiyar1712@gmail.com>

krkn/rollback/config.py

@@ -53,15 +53,30 @@ class SingletonMeta(type):
 class RollbackContent:
     """
     RollbackContent is a dataclass that defines the necessary fields for rollback operations.
+    For cloud-only scenarios (e.g. shut_down) set skip_kubernetes=True and populate
+    cloud_type and instance_ids instead of resource_identifier/namespace.
     """
 
-    resource_identifier: str
+    resource_identifier: str = ""
     namespace: Optional[str] = None
+    cloud_type: Optional[str] = None
+    instance_ids: Optional[tuple] = None
+    skip_kubernetes: bool = False
 
     def __str__(self):
         namespace = f'"{self.namespace}"' if self.namespace else "None"
         resource_identifier = f'"{self.resource_identifier}"'
-        return f"RollbackContent(namespace={namespace}, resource_identifier={resource_identifier})"
+        cloud_type = f'"{self.cloud_type}"' if self.cloud_type else "None"
+        instance_ids = repr(self.instance_ids) if self.instance_ids is not None else "None"
+        return (
+            f"RollbackContent("
+            f"namespace={namespace}, "
+            f"resource_identifier={resource_identifier}, "
+            f"cloud_type={cloud_type}, "
+            f"instance_ids={instance_ids}, "
+            f"skip_kubernetes={self.skip_kubernetes}"
+            f")"
+        )
 
 
 class RollbackContext(str):

krkn/rollback/handler.py

@@ -62,11 +62,12 @@ def set_rollback_context_decorator(func):
 
     def wrapper(self, *args, **kwargs):
         self = cast("AbstractScenarioPlugin", self)
-        # Since `AbstractScenarioPlugin.run_scenarios` will call `self.run` and pass all parameters as `kwargs`
+        # `run_scenarios` passes kwargs; unit tests may call `run` with positional args.
         logger.debug(f"kwargs of ScenarioPlugin.run: {kwargs}")
-        run_uuid = kwargs.get("run_uuid", None)
-        # so we can safely assume that `run_uuid` will be present in `kwargs`
-        assert run_uuid is not None, "run_uuid must be provided in kwargs"
+        run_uuid = kwargs.get("run_uuid")
+        if run_uuid is None and args:
+            run_uuid = args[0]
+        assert run_uuid is not None, "run_uuid must be provided in kwargs or as the first positional argument"
 
         # Set context if run_uuid is available and rollback_handler exists
         if run_uuid and hasattr(self, "rollback_handler"):
@@ -167,7 +168,15 @@ def execute_rollback_version_files(
             rollback_callable, rollback_content = _parse_rollback_module(version_file)
             # Execute the rollback function
             logger.info('Executing rollback callable...')
-            rollback_callable(rollback_content, telemetry_ocp)
+            # Only treat skip_kubernetes as enabled when it is explicitly True.
+            # This avoids accidental truthy values (e.g. Mock objects) disabling telemetry.
+            skip_kubernetes = getattr(rollback_content, "skip_kubernetes", False) is True
+            telemetry_arg = None if skip_kubernetes else telemetry_ocp
+            if telemetry_arg is None and not skip_kubernetes:
+                logger.warning(
+                    "telemetry_ocp is None but skip_kubernetes is not set; rollback callable will receive None"
+                )
+            rollback_callable(rollback_content, telemetry_arg)
             logger.info('Rollback completed.')
             success = True
         except Exception as e:

krkn/rollback/version_template.j2

@@ -3,6 +3,7 @@
 
 from dataclasses import dataclass
 import os
+import sys
 import logging
 from typing import Optional
 
@@ -10,10 +11,22 @@ from krkn_lib.utils import SafeLogger
 from krkn_lib.ocp import KrknOpenshift
 from krkn_lib.telemetry.ocp import KrknTelemetryOpenshift
 
+# Ensure local checkout imports (e.g. `krkn.*`) resolve when executing from a temp dir.
+for candidate in (
+    os.getenv("KRKN_ROOT"),
+    os.getcwd(),
+    os.path.dirname(os.getcwd()),
+):
+    if candidate and os.path.isdir(os.path.join(candidate, "krkn")) and candidate not in sys.path:
+        sys.path.insert(0, candidate)
+
 @dataclass(frozen=True)
 class RollbackContent:
     resource_identifier: str
     namespace: Optional[str] = None
+    cloud_type: Optional[str] = None
+    instance_ids: Optional[tuple] = None
+    skip_kubernetes: bool = False
 
 # Actual rollback callable
 {{ rollback_callable_code }}
@@ -35,16 +48,17 @@ if __name__ == '__main__':
         ]
     )
 
-    # setup logging and get kubeconfig path
-    kubeconfig_path = os.getenv("KUBECONFIG", "~/.kube/config")
     log_directory = os.path.dirname(os.path.abspath(__file__))
     os.makedirs(os.path.join(log_directory, 'logs'), exist_ok=True)
-    # setup SafeLogger for telemetry
-    telemetry_log_path = os.path.join(log_directory, 'logs', 'telemetry.log')
-    safe_logger = SafeLogger(telemetry_log_path)
-    # setup krkn-lib objects
-    lib_openshift = KrknOpenshift(kubeconfig_path=kubeconfig_path)
-    lib_telemetry = KrknTelemetryOpenshift(safe_logger=safe_logger, lib_openshift=lib_openshift)
+    if not rollback_content.skip_kubernetes:
+        # setup logging and get kubeconfig path
+        kubeconfig_path = os.getenv("KUBECONFIG", "~/.kube/config")
+        # setup SafeLogger for telemetry
+        telemetry_log_path = os.path.join(log_directory, 'logs', 'telemetry.log')
+        safe_logger = SafeLogger(telemetry_log_path)
+        # setup krkn-lib objects
+        lib_openshift = KrknOpenshift(kubeconfig_path=kubeconfig_path)
+        lib_telemetry = KrknTelemetryOpenshift(safe_logger=safe_logger, lib_openshift=lib_openshift)
 
     # execute
     logging.info('Executing rollback callable...')

krkn/scenario_plugins/shut_down/shut_down_scenario_plugin.py

@@ -26,12 +26,15 @@ from krkn.scenario_plugins.node_actions.az_node_scenarios import Azure
 from krkn.scenario_plugins.node_actions.gcp_node_scenarios import GCP
 from krkn.scenario_plugins.node_actions.openstack_node_scenarios import OPENSTACKCLOUD
 from krkn.scenario_plugins.node_actions.ibmcloud_node_scenarios import IbmCloud
+from krkn.rollback.handler import set_rollback_context_decorator
+from krkn.rollback.config import RollbackContent
 
 import krkn.scenario_plugins.node_actions.common_node_functions as nodeaction
 
 from krkn_lib.models.k8s import AffectedNodeStatus, AffectedNode
 
 class ShutDownScenarioPlugin(AbstractScenarioPlugin):
+    @set_rollback_context_decorator
     def run(
         self,
         run_uuid: str,
@@ -119,6 +122,19 @@ class ShutDownScenarioPlugin(AbstractScenarioPlugin):
         for _ in range(runs):
             logging.info("Starting cluster_shut_down scenario injection")
             stopping_nodes = set(node_id)
+            
+            # Register rollback callable before shutting down nodes
+            rollback_content = RollbackContent(
+                cloud_type=cloud_type,
+                instance_ids=tuple(node_id),
+                skip_kubernetes=True,
+            )
+            self.rollback_handler.set_rollback_callable(
+                self.rollback_shutdown_nodes,
+                rollback_content
+            )
+            logging.info(f"Registered rollback callable for {len(node_id)} nodes on {cloud_type}")
+            
             self.multiprocess_nodes(cloud_object.stop_instances, node_id, processes)
             stopped_nodes = stopping_nodes.copy()
             start_time = time.time()
@@ -176,3 +192,123 @@ class ShutDownScenarioPlugin(AbstractScenarioPlugin):
 
     def get_scenario_types(self) -> list[str]:
         return ["cluster_shut_down_scenarios"]
+
+    @staticmethod
+    def rollback_shutdown_nodes(rollback_content: RollbackContent, lib_telemetry: KrknTelemetryOpenshift = None):
+        """
+        Rollback function to restore powered-off nodes back to running state.
+        This function works independently of the Kubernetes API server to ensure
+        rollback can function even when all nodes are down (including control plane).
+        
+        :param rollback_content: Rollback content containing node information and cloud provider details.
+        :param lib_telemetry: Instance of KrknTelemetryOpenshift (optional, not used to avoid API server dependency).
+        """
+        import time
+
+        try:
+            # Prefer structured content for cloud-only rollback.
+            cloud_type = rollback_content.cloud_type
+            node_ids = list(rollback_content.instance_ids or ())
+
+            # Backward compatibility for already-serialized rollback files.
+            if not cloud_type or not node_ids:
+                content_parts = rollback_content.resource_identifier.split(":", 1)
+                if len(content_parts) != 2:
+                    logging.error(
+                        f"Invalid rollback content format: {rollback_content.resource_identifier}"
+                    )
+                    return
+                cloud_type = content_parts[0]
+                node_ids_str = content_parts[1]
+                node_ids = [
+                    node_id.strip() for node_id in node_ids_str.split(",") if node_id.strip()
+                ]
+            
+            if not node_ids:
+                logging.warning("No node IDs found in rollback content")
+                return
+            
+            logging.info(f"Rolling back shutdown for {len(node_ids)} nodes on {cloud_type}")
+            logging.info(f"Node IDs: {node_ids}")
+            
+            # Initialize cloud provider
+            # Import at function level to ensure they're available during serialization
+            if cloud_type.lower() == "aws":
+                from krkn.scenario_plugins.node_actions.aws_node_scenarios import AWS
+                cloud_object = AWS()
+            elif cloud_type.lower() == "gcp":
+                from krkn.scenario_plugins.node_actions.gcp_node_scenarios import GCP
+                cloud_object = GCP()
+            elif cloud_type.lower() == "openstack":
+                from krkn.scenario_plugins.node_actions.openstack_node_scenarios import OPENSTACKCLOUD
+                cloud_object = OPENSTACKCLOUD()
+            elif cloud_type.lower() in ["azure", "az"]:
+                from krkn.scenario_plugins.node_actions.az_node_scenarios import Azure
+                cloud_object = Azure()
+            elif cloud_type.lower() in ["ibm", "ibmcloud"]:
+                from krkn.scenario_plugins.node_actions.ibmcloud_node_scenarios import IbmCloud
+                cloud_object = IbmCloud()
+            else:
+                logging.error(f"Unsupported cloud type for rollback: {cloud_type}")
+                return
+            
+            # Start instances one at a time — cloud providers accept a single instance,
+            # not a list (see multiprocess_nodes in cluster_shut_down).
+            logging.info("Starting instances for rollback...")
+            timeout = 300  # 5 minutes timeout for rollback
+            successful_restores = 0
+            failed_restores = []
+
+            for node_id in node_ids:
+                try:
+                    logging.info(f"Starting instance for rollback: {node_id}")
+                    if isinstance(node_id, tuple):
+                        # Azure stores (vm_name, resource_group)
+                        cloud_object.start_instances(node_id[1], node_id[0])
+                    else:
+                        cloud_object.start_instances(node_id)
+                except Exception as e:
+                    logging.error(f"Failed to start instance {node_id}: {e}")
+                    failed_restores.append(node_id)
+                    continue
+
+            logging.info("Waiting for instances to be running...")
+            for node_id in node_ids:
+                if node_id in failed_restores:
+                    continue
+                try:
+                    logging.info(f"Waiting for node {node_id} to be running...")
+                    if isinstance(node_id, tuple):
+                        node_status = cloud_object.wait_until_running(
+                            node_id[1], node_id[0], timeout, None
+                        )
+                    else:
+                        node_status = cloud_object.wait_until_running(node_id, timeout, None)
+                    if node_status:
+                        logging.info(f"Successfully restored node: {node_id}")
+                        successful_restores += 1
+                    else:
+                        logging.warning(f"Timeout waiting for node {node_id} to be running")
+                        failed_restores.append(node_id)
+                except Exception as e:
+                    logging.error(f"Error waiting for node {node_id} to be running: {e}")
+                    failed_restores.append(node_id)
+            
+            # Log rollback summary
+            if successful_restores == len(node_ids):
+                logging.info(f"Rollback completed successfully for all {len(node_ids)} nodes")
+            else:
+                logging.warning(f"Rollback completed with issues: {successful_restores}/{len(node_ids)} nodes restored successfully")
+                if failed_restores:
+                    logging.warning(f"Failed to restore nodes: {failed_restores}")
+            
+            # Wait for cluster components to initialize after rollback
+            if successful_restores > 0:
+                logging.info("Waiting for cluster components to initialize after rollback...")
+                time.sleep(60)  # Shorter wait for rollback scenario
+            
+            logging.info("Rollback of shutdown nodes completed.")
+            
+        except Exception as e:
+            logging.error(f"Failed to rollback shutdown nodes: {e}")
+            raise

tests/test_power_outage_rollback.py

@@ -0,0 +1,258 @@
+"""
+Power Outage Rollback Test Suite
+
+This is the final and comprehensive test suite for the Power Outage Rollback Feature.
+It tests the rollback functionality for restoring powered-off nodes back to running state
+after a power outage scenario fails.
+
+Features tested:
+- Successful node restoration
+- Partial failure handling
+- Invalid content handling
+- Unsupported cloud provider handling
+- Empty node list handling
+- Cloud provider exception handling
+- Multi-cloud provider support (AWS, GCP, Azure, OpenStack, IBM Cloud)
+- Rollback content parsing and edge cases
+"""
+
+import pytest
+import os
+import sys
+from unittest.mock import Mock, patch
+
+# Add the krkn directory to the path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
+
+# Import the actual rollback function and RollbackContent
+from krkn.rollback.config import RollbackContent
+from krkn.scenario_plugins.shut_down.shut_down_scenario_plugin import ShutDownScenarioPlugin
+
+
+class TestPowerOutageRollback:
+    """Test class for Power Outage Rollback functionality."""
+
+    def test_rollback_shutdown_nodes_success(self):
+        """Test successful rollback of shutdown nodes."""
+        rollback_content = RollbackContent(
+            cloud_type="aws",
+            instance_ids=("i-12345", "i-67890"),
+            skip_kubernetes=True,
+        )
+        
+        # Mock telemetry
+        mock_telemetry = Mock()
+        
+        # Mock cloud provider - patch at the import location inside the rollback function
+        with patch('krkn.scenario_plugins.node_actions.aws_node_scenarios.AWS') as mock_aws_class:
+            mock_aws = Mock()
+            mock_aws_class.return_value = mock_aws
+            mock_aws.start_instances.return_value = None
+            mock_aws.wait_until_running.return_value = True
+            
+            # Execute rollback
+            ShutDownScenarioPlugin.rollback_shutdown_nodes(rollback_content, mock_telemetry)
+            
+            # Verify cloud provider methods were called per instance
+            assert mock_aws.start_instances.call_count == 2
+            mock_aws.start_instances.assert_any_call("i-12345")
+            mock_aws.start_instances.assert_any_call("i-67890")
+            assert mock_aws.wait_until_running.call_count == 2
+
+    def test_rollback_shutdown_nodes_partial_failure(self):
+        """Test rollback with partial node restoration failure."""
+        rollback_content = RollbackContent(
+            cloud_type="aws",
+            instance_ids=("i-12345", "i-67890"),
+            skip_kubernetes=True,
+        )
+        
+        # Mock telemetry
+        mock_telemetry = Mock()
+        
+        # Mock cloud provider with partial failure - patch at the import location
+        with patch('krkn.scenario_plugins.node_actions.aws_node_scenarios.AWS') as mock_aws_class:
+            mock_aws = Mock()
+            mock_aws_class.return_value = mock_aws
+            mock_aws.start_instances.return_value = None
+            
+            # First node succeeds, second fails
+            def wait_side_effect(node_id, timeout, affected_node=None):
+                if node_id == "i-12345":
+                    return True
+                return False
+            
+            mock_aws.wait_until_running.side_effect = wait_side_effect
+            
+            # Execute rollback
+            ShutDownScenarioPlugin.rollback_shutdown_nodes(rollback_content, mock_telemetry)
+            
+            # Verify cloud provider methods were called per instance
+            assert mock_aws.start_instances.call_count == 2
+            mock_aws.start_instances.assert_any_call("i-12345")
+            mock_aws.start_instances.assert_any_call("i-67890")
+            assert mock_aws.wait_until_running.call_count == 2
+
+    def test_rollback_shutdown_nodes_invalid_content(self):
+        """Test rollback with invalid rollback content."""
+        # Create invalid rollback content
+        rollback_content = RollbackContent(resource_identifier="invalid_format")
+        
+        # Mock telemetry
+        mock_telemetry = Mock()
+        
+        # Execute rollback - should handle gracefully
+        ShutDownScenarioPlugin.rollback_shutdown_nodes(rollback_content, mock_telemetry)
+        
+        # No cloud provider should be instantiated - patch at the import location
+        with patch('krkn.scenario_plugins.node_actions.aws_node_scenarios.AWS') as mock_aws_class:
+            assert not mock_aws_class.called
+
+    def test_rollback_shutdown_nodes_unsupported_cloud(self):
+        """Test rollback with unsupported cloud provider."""
+        # Create rollback content with unsupported cloud
+        rollback_content = RollbackContent(cloud_type="unsupported", instance_ids=("i-12345",))
+        
+        # Mock telemetry
+        mock_telemetry = Mock()
+        
+        # Execute rollback - should handle gracefully
+        ShutDownScenarioPlugin.rollback_shutdown_nodes(rollback_content, mock_telemetry)
+        
+        # No cloud provider should be instantiated - patch at the import location
+        with patch('krkn.scenario_plugins.node_actions.aws_node_scenarios.AWS') as mock_aws_class:
+            assert not mock_aws_class.called
+
+    def test_rollback_shutdown_nodes_empty_node_list(self):
+        """Test rollback with empty node list."""
+        # Create rollback content with empty node list
+        rollback_content = RollbackContent(cloud_type="aws", instance_ids=tuple())
+        
+        # Mock telemetry
+        mock_telemetry = Mock()
+        
+        # Execute rollback - should handle gracefully
+        ShutDownScenarioPlugin.rollback_shutdown_nodes(rollback_content, mock_telemetry)
+        
+        # No cloud provider should be instantiated - patch at the import location
+        with patch('krkn.scenario_plugins.node_actions.aws_node_scenarios.AWS') as mock_aws_class:
+            assert not mock_aws_class.called
+
+    def test_rollback_shutdown_nodes_cloud_provider_exception(self):
+        """Test rollback when cloud provider start operations fail."""
+        rollback_content = RollbackContent(
+            cloud_type="aws",
+            instance_ids=("i-12345",),
+            skip_kubernetes=True,
+        )
+
+        mock_telemetry = Mock()
+
+        with patch('krkn.scenario_plugins.node_actions.aws_node_scenarios.AWS') as mock_aws_class:
+            mock_aws = Mock()
+            mock_aws_class.return_value = mock_aws
+            mock_aws.start_instances.side_effect = Exception("Cloud API error")
+
+            # Per-node failures are handled gracefully; rollback should not raise.
+            ShutDownScenarioPlugin.rollback_shutdown_nodes(rollback_content, mock_telemetry)
+
+            mock_aws.start_instances.assert_called_once_with("i-12345")
+            mock_aws.wait_until_running.assert_not_called()
+
+    def test_rollback_shutdown_nodes_different_cloud_providers(self):
+        """Test rollback with different cloud providers."""
+        # Mock telemetry
+        mock_telemetry = Mock()
+        
+        cloud_providers = [
+            ("gcp", "krkn.scenario_plugins.node_actions.gcp_node_scenarios.GCP"),
+            ("azure", "krkn.scenario_plugins.node_actions.az_node_scenarios.Azure"),
+            ("openstack", "krkn.scenario_plugins.node_actions.openstack_node_scenarios.OPENSTACKCLOUD"),
+            ("ibm", "krkn.scenario_plugins.node_actions.ibmcloud_node_scenarios.IbmCloud")
+        ]
+        
+        for cloud_type, provider_path in cloud_providers:
+            # Create rollback content
+            rollback_content = RollbackContent(
+                cloud_type=cloud_type,
+                instance_ids=("i-12345",),
+                skip_kubernetes=True,
+            )
+            
+            # Mock cloud provider
+            with patch(provider_path) as mock_provider_class:
+                mock_provider = Mock()
+                mock_provider_class.return_value = mock_provider
+                mock_provider.start_instances.return_value = None
+                mock_provider.wait_until_running.return_value = True
+                
+                # Execute rollback
+                ShutDownScenarioPlugin.rollback_shutdown_nodes(rollback_content, mock_telemetry)
+                
+                # Verify cloud provider methods were called per instance
+                mock_provider.start_instances.assert_called_once_with("i-12345")
+                mock_provider.wait_until_running.assert_called_once()
+
+    def test_rollback_content_parsing(self):
+        """Test structured rollback content fields."""
+        rollback_content = RollbackContent(
+            cloud_type="aws",
+            instance_ids=("i-12345", "i-67890"),
+            skip_kubernetes=True,
+        )
+        assert rollback_content.cloud_type == "aws"
+        assert rollback_content.instance_ids == ("i-12345", "i-67890")
+        assert rollback_content.skip_kubernetes is True
+
+    def test_rollback_content_parsing_edge_cases(self):
+        """Test legacy fallback content parsing with edge cases."""
+        rollback_content = RollbackContent(
+            resource_identifier="aws: i-12345 , i-67890 "
+        )
+        
+        content_parts = rollback_content.resource_identifier.split(":", 1)
+        node_ids = [node_id.strip() for node_id in content_parts[1].split(",") if node_id.strip()]
+        assert node_ids == ["i-12345", "i-67890"]
+        
+        # Test with empty node ID
+        rollback_content = RollbackContent(
+            resource_identifier="aws:i-12345,,i-67890"
+        )
+        
+        content_parts = rollback_content.resource_identifier.split(":", 1)
+        node_ids = [node_id.strip() for node_id in content_parts[1].split(",") if node_id.strip()]
+        assert node_ids == ["i-12345", "i-67890"]
+
+    @patch("krkn.rollback.handler.RollbackConfig.search_rollback_version_files")
+    @patch("krkn.rollback.handler._parse_rollback_module")
+    @patch("os.rename")
+    def test_execute_rollback_passes_none_telemetry_for_skip_kubernetes(
+        self, mock_rename, mock_parse_rollback_module, mock_search_files
+    ):
+        """Cloud-only rollback should receive None telemetry."""
+        from krkn.rollback.handler import execute_rollback_version_files
+
+        version_file = "/tmp/rollback_file.py"
+        mock_search_files.return_value = [version_file]
+        rollback_callable = Mock()
+        rollback_content = RollbackContent(
+            cloud_type="aws",
+            instance_ids=("i-12345",),
+            skip_kubernetes=True,
+        )
+        mock_parse_rollback_module.return_value = (rollback_callable, rollback_content)
+
+        execute_rollback_version_files(
+            telemetry_ocp=Mock(),
+            run_uuid="test-run",
+            scenario_type="cluster_shut_down_scenarios",
+            ignore_auto_rollback_config=True,
+        )
+
+        rollback_callable.assert_called_once_with(rollback_content, None)
+        mock_rename.assert_called_once_with(version_file, f"{version_file}.executed")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])
+

tests/test_rollback.py

@@ -341,7 +341,8 @@ class TestSecureTempDirectories:
                 os.path.expanduser("~"), ".krkn", "rollback"
             )
         assert rollback_versions_dir != "/tmp/kraken-rollback"
-        assert ".krkn/rollback" in rollback_versions_dir
+        normalized = rollback_versions_dir.replace("\\", "/")
+        assert "/.krkn/rollback" in normalized
 
     def test_archive_path_uses_secure_tempdir_when_empty(self):
         """When archive_path is empty, a secure temp directory
@@ -353,7 +354,9 @@ class TestSecureTempDirectories:
             assert os.path.isdir(archive_path)
             assert archive_path != "/tmp"
             mode = oct(os.stat(archive_path).st_mode & 0o777)
-            assert mode == "0o700"
+            # Windows does not reliably expose POSIX permission bits the same way.
+            if os.name != "nt":
+                assert mode == "0o700"
         finally:
             os.rmdir(archive_path)