Refactor inject_args for Direct UUID Support

- Allow passing UUID, list[UUID], and Optional[UUID] directly.
- Refactor conversion logic into helper functions.
- Add tests for UUID, list, and optional UUID handling.

Author: torkashvandmt

orchestrator/utils/state.py

@@ -150,8 +150,14 @@ def _build_arguments(func: StepFunc | InputStepFunc, state: State) -> list:  # n
 
     Raises:
          KeyError: if requested argument is not in the state, or cannot be reconstructed as an initial domain model.
+         ValueError: if requested argument cannot be converted to the expected type.
 
     """
+
+    def _convert_to_uuid(v: Any) -> UUID:
+        """Converts the value to a UUID instance if it is not already one."""
+        return v if isinstance(v, UUID) else UUID(v)
+
     sig = inspect.signature(func)
     arguments: list[Any] = []
     if sig.parameters:
@@ -203,12 +209,24 @@ def _build_arguments(func: StepFunc | InputStepFunc, state: State) -> list:  # n
                 arguments.append(state.get(name, param.default))
             else:
                 try:
-                    arguments.append(state[name])
+                    value = state[name]
+                    if param.annotation == UUID:
+                        arguments.append(_convert_to_uuid(value))
+                    elif is_list_type(param.annotation, UUID):
+                        arguments.append([_convert_to_uuid(item) for item in value])
+                    elif is_optional_type(param.annotation, UUID):
+                        arguments.append(None if value is None else _convert_to_uuid(value))
+                    else:
+                        arguments.append(value)
                 except KeyError as key_error:
                     logger.error("Could not find key in state.", key=name, state=state)
                     raise KeyError(
                         f"Could not find key '{name}' in state. for function {func.__module__}.{func.__qualname__}"
                     ) from key_error
+                except ValueError as value_error:
+                    logger.error("Could not convert value to expected type.", key=name, state=state, value=state[name])
+                    raise ValueError(f"Could not convert value '{state[name]}' to {param.annotation}") from value_error
+
     return arguments
 
 

test/unit_tests/utils/test_state.py

@@ -1,6 +1,6 @@
 # type: ignore
 
-from uuid import uuid4
+from uuid import UUID, uuid4
 
 import pytest
 
@@ -299,3 +299,192 @@ def form_function(generic_sub: GenericProduct) -> State:
     # Test `rt_1` has been persisted to the database with the modifications from the step function.`
     fresh_generic_sub = GenericProduct.from_subscription(state_amended["generic_sub"].subscription_id)
     assert fresh_generic_sub.pb_1.rt_1 is not None
+
+
+def test_uuid_parameter_with_string_uuid():
+    # Test plain UUID parameter: providing a string should result in conversion to a UUID instance.
+
+    @inject_args
+    def step(uuid_param: UUID) -> State:
+        # Return the UUID so we can inspect it
+        return {"result": uuid_param}
+
+    valid_uuid = uuid4()
+    state = {"uuid_param": str(valid_uuid)}
+    new_state = step(state)
+    assert isinstance(new_state["result"], UUID)
+    assert new_state["result"] == valid_uuid
+
+
+def test_uuid_parameter_with_uuid_instance():
+    # Test plain UUID parameter: providing a UUID instance should pass it through unchanged.
+
+    @inject_args
+    def step(uuid_param: UUID) -> State:
+        return {"result": uuid_param}
+
+    valid_uuid = uuid4()
+    state = {"uuid_param": valid_uuid}
+    new_state = step(state)
+    assert isinstance(new_state["result"], UUID)
+    assert new_state["result"] == valid_uuid
+
+
+def test_uuid_parameter_missing_key():
+    # Test that if the key for a UUID parameter is missing, a KeyError is raised.
+
+    @inject_args
+    def step(uuid_param: UUID) -> State:
+        return {"result": uuid_param}
+
+    state = {}  # key missing
+    with pytest.raises(KeyError) as excinfo:
+        step(state)
+
+    assert "Could not find key 'uuid_param' in state. for function" in str(excinfo.value)
+
+
+def test_uuid_parameter_invalid_value():
+    # Test that an invalid UUID string raises a ValueError.
+
+    @inject_args
+    def step(uuid_param: UUID) -> State:
+        return {"result": uuid_param}
+
+    state = {"uuid_param": "not-a-valid-uuid"}
+    with pytest.raises(ValueError) as excinfo:
+        step(state)
+
+    assert f"Could not convert value 'not-a-valid-uuid' to {UUID}" in str(excinfo.value)
+
+
+def test_list_of_uuid_parameter_with_string():
+    # Test list[UUID] parameter: providing a list of valid UUID strings.
+
+    @inject_args
+    def step(uuid_list: list[UUID]) -> State:
+        return {"result": uuid_list}
+
+    valid_uuid = uuid4()
+    state = {"uuid_list": [str(valid_uuid)]}
+    new_state = step(state)
+    result = new_state["result"]
+    assert isinstance(result, list)
+    assert len(result) == 1
+    assert isinstance(result[0], UUID)
+    assert result[0] == valid_uuid
+
+
+def test_list_of_uuid_parameter_with_uuid_instance():
+    # Test list[UUID] parameter: providing a list with a UUID instance.
+
+    @inject_args
+    def step(uuid_list: list[UUID]) -> State:
+        return {"result": uuid_list}
+
+    valid_uuid = uuid4()
+    state = {"uuid_list": [valid_uuid]}
+    new_state = step(state)
+    result = new_state["result"]
+    assert isinstance(result, list)
+    assert len(result) == 1
+    assert isinstance(result[0], UUID)
+    assert result[0] == valid_uuid
+
+
+def test_list_of_uuid_parameter_invalid_value():
+    # Test list[UUID] parameter: an invalid element in the list should cause a ValueError.
+
+    @inject_args
+    def step(uuid_list: list[UUID]) -> State:
+        return {"result": uuid_list}
+
+    state = {"uuid_list": ["invalid-uuid"]}
+    with pytest.raises(ValueError) as excinfo:
+        step(state)
+
+    assert f"Could not convert value '['invalid-uuid']' to {list[UUID]}" in str(excinfo.value)
+
+
+def test_optional_uuid_parameter_with_string():
+    # Test Optional[UUID] parameter: providing a valid UUID string.
+
+    @inject_args
+    def step(optional_uuid: UUID | None) -> State:
+        return {"result": optional_uuid}
+
+    valid_uuid = uuid4()
+    state = {"optional_uuid": str(valid_uuid)}
+    new_state = step(state)
+    result = new_state["result"]
+    assert isinstance(result, UUID)
+    assert result == valid_uuid
+
+
+def test_optional_uuid_parameter_with_uuid_instance():
+    # Test Optional[UUID] parameter: providing a UUID instance.
+
+    @inject_args
+    def step(optional_uuid: UUID | None) -> State:
+        return {"result": optional_uuid}
+
+    valid_uuid = uuid4()
+    state = {"optional_uuid": valid_uuid}
+    new_state = step(state)
+    result = new_state["result"]
+    assert isinstance(result, UUID)
+    assert result == valid_uuid
+
+
+def test_optional_uuid_parameter_with_none():
+    # Test Optional[UUID] parameter: providing None explicitly.
+
+    @inject_args
+    def step(optional_uuid: UUID | None) -> State:
+        return {"result": optional_uuid}
+
+    state = {"optional_uuid": None}
+    new_state = step(state)
+    assert new_state["result"] is None
+
+
+def test_optional_uuid_parameter_missing_key():
+    # Test Optional[UUID] parameter: when the key is missing.
+
+    @inject_args
+    def step(optional_uuid: UUID | None) -> State:
+        return {"result": optional_uuid}
+
+    state = {}  # key missing; no default is provided, so expect a KeyError
+    with pytest.raises(KeyError) as excinfo:
+        step(state)
+
+    assert "Could not find key 'optional_uuid' in state. for function" in str(excinfo.value)
+
+
+def test_optional_uuid_parameter_invalid_value():
+    # Test Optional[UUID] parameter: providing an invalid UUID value should raise ValueError.
+
+    @inject_args
+    def step(optional_uuid: UUID | None) -> State:
+        return {"result": optional_uuid}
+
+    state = {"optional_uuid": "invalid-uuid"}
+    with pytest.raises(ValueError) as excinfo:
+        step(state)
+
+    assert f"Could not convert value 'invalid-uuid' to {UUID | None}" in str(excinfo.value)
+
+
+def test_uuid_parameter_with_default_value():
+    # Test that a default value is used if the key is missing.
+
+    default_uuid = uuid4()
+
+    @inject_args
+    def step(uuid_param: UUID = default_uuid) -> State:
+        return {"result": uuid_param}
+
+    state = {}  # missing key, so the default should be used
+    new_state = step(state)
+    assert new_state["result"] == default_uuid