diff --git a/tests/conftest.py b/tests/conftest.py index 3725363..6fe4d08 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -69,12 +69,27 @@ def node_output_info() -> NodeOutputInfo: @pytest.fixture -def node_requirements_info_data() -> dict: +def node_requirements_info_data( + node_requirements_info_min_data: dict[str, int], +) -> dict[str, bool | int]: """ Some data to define a NodeRequirementsInfo object. """ - return {"cpu": 256, "gpu": True, "memory": 512, "timeout": 60} + return { + "cpu": 256, + "gpu": True, + "memory": 512, + **node_requirements_info_min_data, + } + + +@pytest.fixture +def node_requirements_info_min_data() -> dict[str, int]: + """Minimum data to build a `NodeRequirementsInfo` instance.""" + return { + "timeout": 60, + } @pytest.fixture diff --git a/tests/test_node_info.py b/tests/test_node_info.py index fc33e88..b910067 100644 --- a/tests/test_node_info.py +++ b/tests/test_node_info.py @@ -206,7 +206,28 @@ def test_node_requirements_info(node_requirements_info_data: dict): assert node_requirements_info.model_dump() == node_requirements_info_data -@pytest.mark.parametrize("field", ["cpu", "gpu", "memory", "timeout"]) +def test_node_requirements_info_min( + node_requirements_info_min_data: dict[str, Any], +) -> None: + """ + Asserts that `NodeRequirementsInfo` can be instantiated without + optional values. + """ + instance = NodeRequirementsInfo(**node_requirements_info_min_data) + assert instance.model_dump() == { + "cpu": None, + "gpu": None, + "memory": None, + **node_requirements_info_min_data, + } + + +@pytest.mark.parametrize( + "field", + [ + "timeout", + ], +) def test_node_requirements_info_raise_missing( node_requirements_info_data: dict, field: str ): diff --git a/uncertainty_engine_types/node_info.py b/uncertainty_engine_types/node_info.py index 6ad2ef5..9f68d49 100644 --- a/uncertainty_engine_types/node_info.py +++ b/uncertainty_engine_types/node_info.py @@ -21,9 +21,9 @@ class NodeOutputInfo(BaseModel): class NodeRequirementsInfo(BaseModel): - cpu: int - gpu: bool - memory: int + cpu: int | None = None + gpu: bool | None = None + memory: int | None = None timeout: int