added initial operation context key requirements

Author: Andrei Neagu

services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/generic_scheduler/_core.py

@@ -30,6 +30,7 @@
 from ._errors import (
     CannotCancelWhileWaitingForManualInterventionError,
     NoDataFoundError,
+    OperationInitialContextKeyNotFoundError,
     OperationNotCancellableError,
     StepNameNotInCurrentGroupError,
     StepNotInErrorStateError,
@@ -106,6 +107,12 @@ async def start_operation(
         # check if operation is registered
         operation = OperationRegistry.get_operation(operation_name)
 
+        for required_key in operation.initial_context_required_keys:
+            if required_key not in initial_operation_context:
+                raise OperationInitialContextKeyNotFoundError(
+                    operation_name=operation_name, required_key=required_key
+                )
+
         # NOTE: to ensure reproducibility of operations, the
         # operation steps cannot overwrite keys in the
         # initial context with their results

services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/generic_scheduler/_errors.py

@@ -19,6 +19,12 @@ class OperationNotFoundError(BaseGenericSchedulerError):
     )
 
 
+class OperationInitialContextKeyNotFoundError(BaseGenericSchedulerError):
+    msg_template: str = (
+        "Operation '{operation_name}' required_key='{required_key}' not in initial_operation_context"
+    )
+
+
 class StepNotFoundInoperationError(BaseGenericSchedulerError):
     msg_template: str = (
         "Step '{step_name}' not found steps_names='{steps_names}' for operation '{operation_name}'"

services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/generic_scheduler/_event_after.py

@@ -5,6 +5,7 @@
 from servicelib.logging_utils import log_context
 
 from ._core import start_operation
+from ._errors import OperationInitialContextKeyNotFoundError
 from ._models import (
     EventType,
     OperationContext,
@@ -50,7 +51,12 @@ async def register_to_start_after(
             return
 
         # ensure operation exists
-        OperationRegistry.get_operation(to_start.operation_name)
+        operation = OperationRegistry.get_operation(to_start.operation_name)
+        for required_key in operation.initial_context_required_keys:
+            if required_key not in to_start.initial_context:
+                raise OperationInitialContextKeyNotFoundError(
+                    operation_name=to_start.operation_name, required_key=required_key
+                )
 
         await events_proxy.create_or_update_multiple(
             {

services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/generic_scheduler/_operation.py

@@ -232,9 +232,17 @@ def get_step_subgroup_to_run(self) -> StepsSubGroup:
 
 class Operation:
     def __init__(
-        self, *step_groups: BaseStepGroup, is_cancellable: bool = True
+        self,
+        *step_groups: BaseStepGroup,
+        initial_context_required_keys: set[str] | None = None,
+        is_cancellable: bool = True,
     ) -> None:
         self.step_groups = list(step_groups)
+        self.initial_context_required_keys = (
+            set()
+            if initial_context_required_keys is None
+            else initial_context_required_keys
+        )
         self.is_cancellable = is_cancellable
 
     def __repr__(self) -> str:

services/dynamic-scheduler/tests/unit/services/generic_scheduler/test_generic_scheduler.py

@@ -40,6 +40,12 @@
     register_to_start_after_on_reverted_completed,
     start_operation,
 )
+from simcore_service_dynamic_scheduler.services.generic_scheduler._core import (
+    OperationContext,
+)
+from simcore_service_dynamic_scheduler.services.generic_scheduler._errors import (
+    OperationInitialContextKeyNotFoundError,
+)
 from utils import (
     BaseExpectedStepOrder,
     ExecuteRandom,
@@ -438,8 +444,8 @@ async def test_can_recover_from_interruption(
     ],
 )
 async def test_run_operation_after(
-    app: FastAPI,
     preserve_caplog_for_async_logging: None,
+    app: FastAPI,
     steps_call_order: list[tuple[str, str]],
     register_operation: Callable[[OperationName, Operation], None],
     register_at_creation: bool,
@@ -481,3 +487,67 @@ async def test_run_operation_after(
 
     await ensure_expected_order(steps_call_order, expected_order)
     await ensure_keys_in_store(app, expected_keys=set())
+
+
+async def test_missing_initial_context_key_from_operation(
+    preserve_caplog_for_async_logging: None,
+    app: FastAPI,
+    register_operation: Callable[[OperationName, Operation], None],
+):
+    good_operation_name: OperationName = "good"
+    bad_operation_name: OperationName = "bad"
+
+    operation = Operation(
+        SingleStepGroup(_ShortSleep), initial_context_required_keys={"required_key"}
+    )
+    register_operation(good_operation_name, operation)
+    register_operation(bad_operation_name, operation)
+
+    common_initial_context = {"unsued1": "value1", "unsued2": "value2"}
+    good_initial_context: OperationContext = {
+        "required_key": "soeme_value",
+        **common_initial_context,
+    }
+    bad_initial_context: OperationContext = {**common_initial_context}
+
+    bad_operation_to_start = OperationToStart(
+        operation_name=bad_operation_name, initial_context=bad_initial_context
+    )
+
+    # 1. check it works
+    await start_operation(app, bad_operation_name, good_initial_context)
+
+    # 2. check it raises with a bad context
+    with pytest.raises(OperationInitialContextKeyNotFoundError):
+        await start_operation(app, bad_operation_name, bad_initial_context)
+
+    with pytest.raises(OperationInitialContextKeyNotFoundError):
+        await start_operation(
+            app,
+            good_operation_name,
+            good_initial_context,
+            on_execute_completed=bad_operation_to_start,
+            on_revert_completed=None,
+        )
+
+    with pytest.raises(OperationInitialContextKeyNotFoundError):
+        await start_operation(
+            app,
+            good_operation_name,
+            good_initial_context,
+            on_execute_completed=None,
+            on_revert_completed=bad_operation_to_start,
+        )
+
+    # 3. register_to_start_after... raises with a bad context
+    schedule_id = await start_operation(app, bad_operation_name, good_initial_context)
+
+    with pytest.raises(OperationInitialContextKeyNotFoundError):
+        await register_to_start_after_on_executed_completed(
+            app, schedule_id, to_start=bad_operation_to_start
+        )
+
+    with pytest.raises(OperationInitialContextKeyNotFoundError):
+        await register_to_start_after_on_reverted_completed(
+            app, schedule_id, to_start=bad_operation_to_start
+        )