Use a status enum instead of hard-coded strings for step status

this allow separation of "Skipped" into skipped by error and by user

Author: neomatamune

cosmotech/orchestrator/api/run.py

@@ -21,7 +21,7 @@
 
 from cosmotech.orchestrator import VERSION
 from cosmotech.orchestrator.core.orchestrator import Orchestrator
-from cosmotech.orchestrator.core.step import Step
+from cosmotech.orchestrator.core.step import Step, StepStatus
 from cosmotech.orchestrator.utils.logger import LOGGER
 from cosmotech.orchestrator.utils.translate import T
 
@@ -144,7 +144,7 @@ def run_template(
             LOGGER.info(v[0].simple_repr())
             LOGGER.debug(str(v[0]))
             results[k] = v[0]
-            if v[0].status == "RunError":
+            if v[0].status == StepStatus.ERROR:
                 success = False
 
         if exit_handlers:

cosmotech/orchestrator/core/orchestrator.py

@@ -13,6 +13,7 @@
 
 from cosmotech.orchestrator.core.runner import Runner
 from cosmotech.orchestrator.core.step import Step
+from cosmotech.orchestrator.core.step import StepStatus
 from cosmotech.orchestrator.templates.library import Library
 from cosmotech.orchestrator.templates.plugin import Plugin
 from cosmotech.orchestrator.utils.logger import LOGGER
@@ -102,7 +103,7 @@ def _load_from_json_content(
             for _precedent in _step.precedents:
                 if isinstance(_precedent, str):
                     if _precedent not in _steps:
-                        _step.status = "Error"
+                        _step.status = StepStatus.ERROR
                         raise ValueError(
                             T("csm-orc.orchestrator.core.orchestrator.step_not_exists").format(step_id=_precedent)
                         )

cosmotech/orchestrator/core/step.py

@@ -26,6 +26,19 @@
 from cosmotech.orchestrator.utils.translate import T
 
 
+from enum import Enum
+
+
+class StepStatus(Enum):
+    CREATED = 0
+    INITIALIZED = 1
+    SUCCESS = 2
+    SKIPPED_BY_USER = 3
+    SKIPPED_AFTER_FAILURE = 4
+    ERROR = 5
+    DRY_RUN = 6
+
+
 @dataclass
 class Step:
     id: str = field()
@@ -40,7 +53,7 @@ class Step:
     inputs: dict = field(default_factory=dict)
     captured_output: dict = field(default_factory=dict)
     loaded = False
-    status = None
+    status: StepStatus = StepStatus.CREATED
     skipped = False
     stop_library_load: InitVar[bool] = field(default=False, repr=False)
 
@@ -94,7 +107,7 @@ def __load_command_from_library(self):
         self.display_command_id = self.commandId
         command: CommandTemplate = library.find_template_by_name(self.commandId)
         if command is None:
-            self.status = "Error"
+            self.StepStatus = StepStatus.ERROR
             LOGGER.error(
                 T("csm-orc.orchestrator.core.step.template_not_found").format(
                     step_id=self.display_id, command_id=self.display_command_id
@@ -119,13 +132,13 @@ def __load_command_from_library(self):
 
     def __post_init__(self, stop_library_load):
         if not bool(self.command) ^ bool(self.commandId):
-            self.status = "Error"
+            self.status = StepStatus.ERROR
             raise ValueError(T("csm-orc.orchestrator.core.step.command_required"))
         tmp_env = dict()
         for k, v in self.environment.items():
             tmp_env[k] = EnvironmentVariable(k, **v)
         self.environment = tmp_env
-        self.status = "Init"
+        self.status = StepStatus.INITIALIZED
         self.display_id = self.id
         if self.commandId and not stop_library_load:
             self.__load_command_from_library()
@@ -185,26 +198,30 @@ def run(self, dry: bool = False, previous=None, input_data: dict = None, as_exit
             step_type = "exit handler"
 
         LOGGER.info(T("csm-orc.orchestrator.core.step.starting").format(step_type=step_type, step_id=self.display_id))
-        self.status = "Ready"
 
-        if isinstance(previous, dict) and any(map(lambda a: a not in ["Done", "DryRun"], previous.values())):
+        if isinstance(previous, dict) and any(
+            map(
+                lambda a: a not in [StepStatus.SUCCESS, StepStatus.DRY_RUN, StepStatus.SKIPPED_BY_USER],
+                previous.values(),
+            )
+        ):
             LOGGER.warning(
                 T("csm-orc.orchestrator.core.step.skipping_previous_errors").format(
                     step_type=step_type, step_id=self.display_id
                 )
             )
-            self.status = "Skipped"
+            self.status = StepStatus.SKIPPED_AFTER_FAILURE
 
-        if self.status == "Ready":
+        if self.status == StepStatus.INITIALIZED:
             if self.skipped:
                 LOGGER.info(
                     T("csm-orc.orchestrator.core.step.skipping_as_required").format(
                         step_type=step_type, step_id=self.display_id
                     )
                 )
-                self.status = "Done"
+                self.status = StepStatus.SKIPPED_BY_USER
             elif dry:
-                self.status = "DryRun"
+                self.status = StepStatus.DRY_RUN
             else:
                 # Set up environment with input data
                 _e = self._effective_env()
@@ -355,7 +372,7 @@ def run(self, dry: bool = False, previous=None, input_data: dict = None, as_exit
                             step_type=step_type, step_id=self.display_id
                         )
                     )
-                    self.status = "Done"
+                    self.status = StepStatus.SUCCESS
 
                 except subprocess.CalledProcessError as e:
                     LOGGER.error(
@@ -364,7 +381,7 @@ def run(self, dry: bool = False, previous=None, input_data: dict = None, as_exit
                         )
                     )
                     LOGGER.error(str(e))
-                    self.status = "RunError"
+                    self.status = StepStatus.ERROR
 
         return self.status
 
@@ -379,9 +396,9 @@ def check_env(self):
     def simple_repr(self):
         if self.description:
             return T("csm-orc.orchestrator.core.step.info.simple_repr").format(
-                id=self.id, status=self.status, description=self.description
+                id=self.id, status=self.status.name, description=self.description
             )
-        return T("csm-orc.orchestrator.core.step.info.simple_repr_no_desc").format(id=self.id, status=self.status)
+        return T("csm-orc.orchestrator.core.step.info.simple_repr_no_desc").format(id=self.id, status=self.status.name)
 
     def __str__(self):
         r = list()
@@ -406,5 +423,5 @@ def __str__(self):
             r.append(T("csm-orc.orchestrator.core.step.info.use_system_env"))
         if self.skipped:
             r.append(T("csm-orc.orchestrator.core.step.info.skipped"))
-        r.append(T("csm-orc.orchestrator.core.step.info.status").format(status=self.status))
+        r.append(T("csm-orc.orchestrator.core.step.info.status").format(status=self.status.name))
         return "\n".join(r)

tests/unit/orchestrator/api/test_run.py

@@ -8,6 +8,7 @@
 from cosmotech.orchestrator.api.run import generate_env_file
 from cosmotech.orchestrator.api.run import run_template
 from cosmotech.orchestrator.api.run import validate_template
+from cosmotech.orchestrator.core.step import StepStatus
 
 
 class TestValidateTemplate:
@@ -151,11 +152,11 @@ def test_successful_run_with_all_steps_done(self, mock_orchestrator_class):
 
         # Create mock steps and graph
         mock_step1 = MagicMock()
-        mock_step1.status = "Done"
+        mock_step1.status = StepStatus.SUCCESS
         mock_step1.simple_repr.return_value = "Step 1 Done"
 
         mock_step2 = MagicMock()
-        mock_step2.status = "Done"
+        mock_step2.status = StepStatus.SUCCESS
         mock_step2.simple_repr.return_value = "Step 2 Done"
 
         mock_graph = MagicMock()
@@ -185,11 +186,11 @@ def test_run_with_error_step(self, mock_orchestrator_class):
 
         # Create mock steps and graph
         mock_step1 = MagicMock()
-        mock_step1.status = "Done"
+        mock_step1.status = StepStatus.SUCCESS
         mock_step1.simple_repr.return_value = "Step 1 Done"
 
         mock_step2 = MagicMock()
-        mock_step2.status = "RunError"
+        mock_step2.status = StepStatus.ERROR
         mock_step2.simple_repr.return_value = "Step 2 Error"
 
         mock_graph = MagicMock()
@@ -221,7 +222,7 @@ def test_run_with_exit_handlers(self, mock_step_class, mock_library_class, mock_
 
         # Create mock steps and graph
         mock_step1 = MagicMock()
-        mock_step1.status = "Done"
+        mock_step1.status = StepStatus.SUCCESS
         mock_step1.simple_repr.return_value = "Step 1 Done"
 
         mock_graph = MagicMock()
@@ -262,7 +263,7 @@ def test_run_with_dry_run_flag(self, mock_orchestrator_class):
 
         # Create mock steps and graph
         mock_step1 = MagicMock()
-        mock_step1.status = "DryRun"
+        mock_step1.status = StepStatus.DRY_RUN
         mock_step1.simple_repr.return_value = "Step 1 DryRun"
 
         mock_graph = MagicMock()
@@ -287,7 +288,7 @@ def test_run_with_display_env_flag(self, mock_orchestrator_class):
 
         # Create mock steps and graph
         mock_step1 = MagicMock()
-        mock_step1.status = "Done"
+        mock_step1.status = StepStatus.SUCCESS
         mock_step1.simple_repr.return_value = "Step 1 Done"
 
         mock_graph = MagicMock()
@@ -312,11 +313,11 @@ def test_run_with_skipped_steps(self, mock_orchestrator_class):
 
         # Create mock steps and graph
         mock_step1 = MagicMock()
-        mock_step1.status = "Done"
+        mock_step1.status = StepStatus.SUCCESS
         mock_step1.simple_repr.return_value = "Step 1 Done"
 
         mock_step2 = MagicMock()
-        mock_step2.status = "Skipped"
+        mock_step2.status = StepStatus.SKIPPED_BY_USER
         mock_step2.simple_repr.return_value = "Step 2 Skipped"
 
         mock_graph = MagicMock()

tests/unit/orchestrator/core/test_step.py

@@ -8,7 +8,7 @@
 import sys
 from pathlib import Path
 
-from cosmotech.orchestrator.core.step import Step
+from cosmotech.orchestrator.core.step import Step, StepStatus
 from cosmotech.orchestrator.core.environment import EnvironmentVariable
 from cosmotech.orchestrator.templates.library import Library
 
@@ -203,8 +203,8 @@ def test_run_successful_command(self, mock_remove, mock_popen, mock_temp_file):
         result = step.run()
 
         # Verify
-        assert result == "Done"
-        assert step.status == "Done"
+        assert result == StepStatus.SUCCESS
+        assert step.status == StepStatus.SUCCESS
         assert "output1" in step.captured_output
         assert step.captured_output["output1"] == "value1"
         mock_temp_file.assert_called_once()
@@ -233,8 +233,8 @@ def test_run_command_with_error(self, mock_remove, mock_popen, mock_temp_file):
         result = step.run()
 
         # Verify
-        assert result == "RunError"
-        assert step.status == "RunError"
+        assert result == StepStatus.ERROR
+        assert step.status == StepStatus.ERROR
         mock_temp_file.assert_called_once()
         mock_popen.assert_called_once()
         mock_remove.assert_called_once_with("/tmp/test_file")
@@ -247,8 +247,8 @@ def test_run_with_dry_run(self):
         result = step.run(dry=True)
 
         # Verify
-        assert result == "DryRun"
-        assert step.status == "DryRun"
+        assert result == StepStatus.DRY_RUN
+        assert step.status == StepStatus.DRY_RUN
 
     def test_run_with_skipped_step(self):
         # Setup
@@ -259,19 +259,19 @@ def test_run_with_skipped_step(self):
         result = step.run()
 
         # Verify
-        assert result == "Done"
-        assert step.status == "Done"
+        assert result == StepStatus.SKIPPED_BY_USER
+        assert step.status == StepStatus.SKIPPED_BY_USER
 
     def test_run_with_previous_errors(self):
         # Setup
         step = Step(id="test-step", command="echo", arguments=["Hello", "World"])
 
         # Execute
-        result = step.run(previous={"step1": "RunError"})
+        result = step.run(previous={"step1": StepStatus.ERROR})
 
         # Verify
-        assert result == "Skipped"
-        assert step.status == "Skipped"
+        assert result == StepStatus.SKIPPED_AFTER_FAILURE
+        assert step.status == StepStatus.SKIPPED_AFTER_FAILURE
 
     @patch("tempfile.NamedTemporaryFile")
     @patch("subprocess.Popen")
@@ -300,7 +300,7 @@ def test_run_with_input_data(self, mock_remove, mock_popen, mock_temp_file):
         result = step.run(input_data={"input1": "input_value1"})
 
         # Verify
-        assert result == "Done"
+        assert result == StepStatus.SUCCESS
         mock_popen.assert_called_once()
         # Check that environment variables were set correctly
         env = mock_popen.call_args[1]["env"]
@@ -352,7 +352,7 @@ def test_run_with_default_input_value(self, mock_remove, mock_popen, mock_temp_f
         result = step.run(input_data={})
 
         # Verify
-        assert result == "Done"
+        assert result == StepStatus.SUCCESS
         # Check that environment variables were set correctly
         env = mock_popen.call_args[1]["env"]
         assert "INPUT_VAR1" in env
@@ -397,14 +397,14 @@ def test_check_env_returns_empty_dict_for_skipped_step(self):
     def test_simple_repr(self):
         # Setup
         step = Step(id="test-step", command="echo", description="Test step")
-        step.status = "Done"
+        step.status = StepStatus.SUCCESS
 
         # Execute
         result = step.simple_repr()
 
         # Verify
         assert "test-step" in result
-        assert "Done" in result
+        assert StepStatus.SUCCESS.name in result
         assert "Test step" in result
 
     def test_str_representation(self):
@@ -417,7 +417,7 @@ def test_str_representation(self):
             environment={"TEST_VAR": {"value": "test_value", "description": "Test var"}},
             useSystemEnvironment=True,
         )
-        step.status = "Done"
+        step.status = StepStatus.SUCCESS
 
         # Execute
         result = str(step)
@@ -428,7 +428,7 @@ def test_str_representation(self):
         assert "Test step" in result
         assert "TEST_VAR" in result
         assert "Test var" in result
-        assert "Done" in result
+        assert StepStatus.SUCCESS.name in result
 
 
 class TestOutputParser: