Improve Step validation, added unit test for the input

Author: GioeleB00

src/app/config/constants.py

@@ -69,7 +69,7 @@ class Distribution(StrEnum):
 # are accepted by the Pydantic schema.
 
 
-class EndpointIO(StrEnum):
+class EndpointStepIO(StrEnum):
     """
     I/O-bound operation categories that can occur inside an endpoint *step*.
 
@@ -100,7 +100,7 @@ class EndpointIO(StrEnum):
     CACHE      = "io_cache"
 
 
-class EndpointCPU(StrEnum):
+class EndpointStepCPU(StrEnum):
     """
     CPU-bound operation categories inside an endpoint step.
 
@@ -112,7 +112,7 @@ class EndpointCPU(StrEnum):
     CPU_BOUND_OPERATION  = "cpu_bound_operation"
 
 
-class EndpointRAM(StrEnum):
+class EndpointStepRAM(StrEnum):
     """
     Memory-related operations inside a step.
 
@@ -123,7 +123,7 @@ class EndpointRAM(StrEnum):
     RAM = "ram"
 
 
-class MetricKeys(StrEnum):
+class Metrics(StrEnum):
     """
     Keys used inside the ``metrics`` dictionary of a *step*.
 
@@ -137,6 +137,7 @@ class MetricKeys(StrEnum):
 
     NETWORK_LATENCY = "network_latency"
     CPU_TIME        = "cpu_time"
+    IO_WAITING_TIME = "io_waiting_time"
     NECESSARY_RAM   = "necessary_ram"
 
 # ======================================================================

src/app/schemas/system_topology_schema/endpoint_schema.py

@@ -1,30 +1,86 @@
 """Defining the input schema for the requests handler"""
 
-from pydantic import BaseModel, PositiveFloat, PositiveInt, field_validator
+from pydantic import (
+    BaseModel,
+    PositiveFloat,
+    PositiveInt,
+    field_validator,
+    model_validator,
+)
 
-from app.config.constants import EndpointCPU, EndpointIO, EndpointRAM, MetricKeys
+from app.config.constants import (
+    EndpointStepCPU,
+    EndpointStepIO,
+    EndpointStepRAM,
+    Metrics,
+)
 
 
 class Step(BaseModel):
-    """Full step structure to be validated with pydantic"""
+    """
+    Steps to be executed inside an endpoint in terms of
+    the resources needed to accomplish the single step
+    """
 
-    # TODO ADD validation to couple kind and metrics only when they are
-    # valid, for example ram cannot have cpu time and so on
+    kind: EndpointStepIO | EndpointStepCPU | EndpointStepRAM
+    step_metrics: dict[Metrics, PositiveFloat | PositiveInt]
 
-    kind: EndpointIO | EndpointCPU | EndpointRAM
-    metrics: dict[MetricKeys, PositiveFloat | PositiveInt]
+    @field_validator("step_metrics", mode="before")
+    def ensure_non_empty(
+        cls, # noqa: N805
+        v: dict[Metrics, PositiveFloat | PositiveInt],
+        ) -> dict[Metrics, PositiveFloat | PositiveInt]:
+        """Ensure the dict step metrics exist"""
+        if not v:
+            msg = "step_metrics cannot be empty"
+            raise ValueError(msg)
+        return v
 
-    @field_validator("metrics", mode="before")
-    def ensure_metric_exist_positive(
+    @model_validator(mode="after") # type: ignore[arg-type]
+    def ensure_coherence_kind_metrics(
         cls, # noqa: N805
-        v:  dict[MetricKeys, float | int],
-        ) -> dict[MetricKeys, float | int]:
-        """Ensure the measure of an operation exist and is positive"""
-        for key, value in v.items():
-            if not value:
-                msg = f"{key} must be a positive number"
+        model: "Step",
+        ) -> "Step":
+        """
+        Validation to couple kind and metrics only when they are
+        valid for example ram cannot have associated a cpu time
+        """
+        metrics_keys = set(model.step_metrics)
+
+        # Control of the length of the set to be sure only on key is passed
+        if len(metrics_keys) != 1:
+            msg = "step_metrics must contain exactly one entry"
+            raise ValueError(msg)
+
+        # Coherence CPU bound operation and metric
+        if isinstance(model.kind, EndpointStepCPU):
+            if metrics_keys != {Metrics.CPU_TIME}:
+                msg = (
+                        "The metric to quantify a CPU BOUND step"
+                        f"must be {Metrics.CPU_TIME}"
+                    )
                 raise ValueError(msg)
-        return v
+
+        # Coherence RAM operation and metric
+        elif isinstance(model.kind, EndpointStepRAM):
+            if metrics_keys != {Metrics.NECESSARY_RAM}:
+                msg = (
+                       "The metric to quantify a RAM step"
+                       f"must be {Metrics.NECESSARY_RAM}"
+                    )
+                raise ValueError(msg)
+
+        # Coherence I/O operation and metric
+        elif metrics_keys != {Metrics.IO_WAITING_TIME}:
+            msg = (
+                "The metric to quantify an I/O step"
+                f"must be {Metrics.IO_WAITING_TIME}"
+            )
+            raise ValueError(msg)
+
+        return model
+
+
 
 
 class Endpoint(BaseModel):

src/app/schemas/system_topology_schema/full_system_topology_schema.py

@@ -8,6 +8,7 @@
 
 from pydantic import (
     BaseModel,
+    ConfigDict,
     Field,
     PositiveInt,
     field_validator,
@@ -130,10 +131,7 @@ def unique_ids(
             raise ValueError(msg)
         return model
 
-    class Config:
-        """strict control over the acceptable keys"""
-
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
 
 #-------------------------------------------------------------

tests/unit/input_sructure/test_endpoint_input.py

@@ -0,0 +1,125 @@
+"""Unit tests for the Endpoint and Step Pydantic schemas."""
+
+from __future__ import annotations
+
+import pytest
+from pydantic import ValidationError
+
+from app.config.constants import (
+    EndpointStepCPU,
+    EndpointStepIO,
+    EndpointStepRAM,
+    Metrics,
+)
+from app.schemas.system_topology_schema.endpoint_schema import Endpoint, Step
+
+
+# --------------------------------------------------------------------------- #
+# Helper functions to build minimal valid Step objects
+# --------------------------------------------------------------------------- #
+def cpu_step(value: float = 0.1) -> Step:
+    """Return a minimal valid CPU-bound Step."""
+    return Step(
+        kind=EndpointStepCPU.CPU_BOUND_OPERATION,
+        step_metrics={Metrics.CPU_TIME: value},
+    )
+
+
+def ram_step(value: int = 128) -> Step:
+    """Return a minimal valid RAM Step."""
+    return Step(
+        kind=EndpointStepRAM.RAM,
+        step_metrics={Metrics.NECESSARY_RAM: value},
+    )
+
+
+def io_step(value: float = 0.05) -> Step:
+    """Return a minimal valid I/O Step."""
+    return Step(
+        kind=EndpointStepIO.WAIT,
+        step_metrics={Metrics.IO_WAITING_TIME: value},
+    )
+
+
+# --------------------------------------------------------------------------- #
+# Positive test cases
+# --------------------------------------------------------------------------- #
+def test_valid_cpu_step() -> None:
+    """Test that a CPU step with correct 'cpu_time' metric passes validation."""
+    step = cpu_step()
+    # The metric value must match the input
+    assert step.step_metrics[Metrics.CPU_TIME] == 0.1
+
+
+def test_valid_ram_step() -> None:
+    """Test that a RAM step with correct 'necessary_ram' metric passes validation."""
+    step = ram_step()
+    assert step.step_metrics[Metrics.NECESSARY_RAM] == 128
+
+
+def test_valid_io_step() -> None:
+    """Test that an I/O step with correct 'io_waiting_time' metric passes validation."""
+    step = io_step()
+    assert step.step_metrics[Metrics.IO_WAITING_TIME] == 0.05
+
+
+def test_endpoint_with_mixed_steps() -> None:
+    """Test that an Endpoint with multiple valid Step instances normalizes the name."""
+    ep = Endpoint(
+        endpoint_name="/Predict",
+        steps=[cpu_step(), ram_step(), io_step()],
+    )
+    # endpoint_name should be lowercased by the validator
+    assert ep.endpoint_name == "/predict"
+    # All steps should be present in the list
+    assert len(ep.steps) == 3
+
+
+# --------------------------------------------------------------------------- #
+# Negative test cases
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize(
+    ("kind", "bad_metrics"),
+    [
+        # CPU step with RAM metric
+        (EndpointStepCPU.CPU_BOUND_OPERATION, {Metrics.NECESSARY_RAM: 64}),
+        # RAM step with CPU metric
+        (EndpointStepRAM.RAM, {Metrics.CPU_TIME: 0.2}),
+        # I/O step with CPU metric
+        (EndpointStepIO.DB, {Metrics.CPU_TIME: 0.05}),
+    ],
+)
+def test_incoherent_kind_metric_pair(
+    kind: EndpointStepCPU | EndpointStepRAM | EndpointStepIO,
+    bad_metrics: dict[Metrics, float | int],
+) -> None:
+    """Test that mismatched kind and metric combinations raise ValidationError."""
+    with pytest.raises(ValidationError):
+        Step(kind=kind, step_metrics=bad_metrics)
+
+
+def test_multiple_metrics_not_allowed() -> None:
+    """Test that providing multiple metrics in a single Step raises ValidationError."""
+    with pytest.raises(ValidationError):
+        Step(
+            kind=EndpointStepCPU.CPU_BOUND_OPERATION,
+            step_metrics={
+                Metrics.CPU_TIME: 0.1,
+                Metrics.NECESSARY_RAM: 64,
+            },
+        )
+
+
+def test_empty_metrics_rejected() -> None:
+    """Test that an empty metrics dict is rejected by the validator."""
+    with pytest.raises(ValidationError):
+        Step(kind=EndpointStepCPU.CPU_BOUND_OPERATION, step_metrics={})
+
+
+def test_wrong_metric_name_for_io() -> None:
+    """Test that an I/O step with a non-I/O metric key is rejected."""
+    with pytest.raises(ValidationError):
+        Step(
+            kind=EndpointStepIO.CACHE,
+            step_metrics={Metrics.NECESSARY_RAM: 64},
+        )