Handle access to undefined states gracefully

Fixes: #76
Signed-off-by: Tobias Wolf <[email protected]>

Author: NotTheEvilOne

src/rookify/modules/create_rook_cluster/main.py

@@ -16,9 +16,11 @@ class CreateRookClusterHandler(ModuleHandler):
     ]
 
     def preflight(self) -> None:
-        if self.machine.get_execution_state_data(
-            "CreateRookClusterHandler", "generated", default_value=False
-        ):
+        cluster_definition = self.machine.get_preflight_state_data(
+            "CreateRookClusterHandler", "cluster_definition"
+        )
+
+        if cluster_definition is not None:
             return
 
         state_data = self.machine.get_preflight_state("AnalyzeCephHandler").data
@@ -64,13 +66,9 @@ def preflight(self) -> None:
             if mgr_count > 5:
                 mgr_count = rook_config["cluster"].get("max_mgr_count", 5)
 
-            self.machine.get_execution_state(
-                "CreateRookClusterHandler"
-            ).mgr_count = mgr_count
-
-            self.machine.get_execution_state(
-                "CreateRookClusterHandler"
-            ).mon_count = mon_count
+            state = self.machine.get_preflight_state("CreateRookClusterHandler")
+            state.mgr_count = mgr_count
+            state.mon_count = mon_count
 
             cluster_definition_values = {
                 "cluster_name": rook_config["cluster"]["name"],

src/rookify/modules/machine.py

@@ -83,15 +83,25 @@ def _get_state_tags_data(self, name: str) -> Dict[str, Any]:
         return data
 
     def get_execution_state(self, name: str) -> Any:
-        return self.get_state(self.__class__.STATE_NAME_EXECUTION_PREFIX + name)
+        state_name = self.__class__.STATE_NAME_EXECUTION_PREFIX + name
+
+        if state_name not in self.states:
+            return None
+
+        return self.get_state(state_name)
 
     def get_execution_state_data(
         self, name: str, tag: str, default_value: Any = None
     ) -> Any:
         return getattr(self.get_execution_state(name), tag, default_value)
 
     def get_preflight_state(self, name: str) -> Any:
-        return self.get_state(self.__class__.STATE_NAME_PREFLIGHT_PREFIX + name)
+        state_name = self.__class__.STATE_NAME_PREFLIGHT_PREFIX + name
+
+        if state_name not in self.states:
+            return None
+
+        return self.get_state(state_name)
 
     def get_preflight_state_data(
         self, name: str, tag: str, default_value: Any = None

src/rookify/modules/migrate_mds_pools/main.py

@@ -40,9 +40,9 @@ def preflight(self) -> None:
                     if pool_data_osd_name not in migrated_pools:
                         migrated_pools.append(pool_data_osd_name)
 
-                self.machine.get_execution_state(
-                    "MigrateMdsPoolsHandler"
-                ).migrated_pools = migrated_pools
+                state = self.machine.get_execution_state("MigrateMdsPoolsHandler")
+                if state is not None:
+                    state.migrated_pools = migrated_pools
 
                 continue
 

src/rookify/modules/migrate_mgrs/main.py

@@ -90,7 +90,7 @@ def _migrate_mgr(self, mgr_host: str) -> None:
             "MigrateMgrsHandler"
         ).migrated_mgrs = migrated_mgrs
 
-        mgr_count_expected = self.machine.get_execution_state_data(
+        mgr_count_expected = self.machine.get_preflight_state_data(
             "CreateRookClusterHandler", "mgr_count", default_value=3
         )
 

src/rookify/modules/migrate_mons/main.py

@@ -106,7 +106,7 @@ def _migrate_mon(self, mon: Dict[str, Any]) -> None:
             "MigrateMonsHandler",
         ).migrated_mons = migrated_mons
 
-        mon_count_expected = self.machine.get_execution_state_data(
+        mon_count_expected = self.machine.get_preflight_state_data(
             "CreateRookClusterHandler", "mon_count", default_value=3
         )