Merge pull request #1017 from PowerGridModel/feature/make-test-dataset-helper-fn

Test case creation: add helper function to create test case

Author: nitbharambe

src/power_grid_model/utils.py

@@ -48,6 +48,12 @@
 _DEPRECATED_JSON_DESERIALIZATION_MSG = f"{_DEPRECATED_FUNCTION_MSG} Please use json_deserialize_to_file instead."
 _DEPRECATED_JSON_SERIALIZATION_MSG = f"{_DEPRECATED_FUNCTION_MSG} Please use json_serialize_from_file instead."
 
+LICENSE_TEXT = (
+    "SPDX-FileCopyrightText: Contributors to the Power Grid Model project <[email protected]>\n\n"
+    "SPDX-License-Identifier: MPL-2.0"
+    "\n"
+)
+
 
 def get_dataset_scenario(dataset: BatchDataset, scenario: int) -> SingleDataset:
     """
@@ -398,3 +404,62 @@ def self_test():
             print("Self test finished.")
         except Exception as e:
             raise PowerGridError from e
+
+
+def _make_test_case(
+    *,
+    output_path: Path,
+    input_data: SingleDataset,
+    params: dict,
+    output_data: Dataset,
+    output_dataset_type: DatasetType,
+    update_data: Dataset | None = None,
+):
+    """
+    Create and save a validation test case dataset, including input, update (optional), output, and parameters.
+
+    Args:
+        save_path: Directory path where the test case files will be saved.
+        input_data: Input dataset for the test case.
+        params: Dictionary of parameters used for the test case. It may include calculation method, tolerances, etc.
+            An example of parameters could be:
+            params = {
+                "calculation_method": "newton_raphson",
+                "rtol": 1e-6,
+                "atol": 1e-6,
+            }
+        output_data: Output dataset for the test case.
+        output_dataset_type: The type of the output dataset (e.g., sym_output, asym_output, sc_output).
+        update_data: Optional batch update dataset.
+
+    Raises:
+        ValueError: If the output_dataset_type is not recognized.
+
+    Side Effects:
+        Writes JSON files for input, update (if provided), output, and parameters,
+        all relevant license files, to save_path.
+    """
+    output_file_stem = output_dataset_type.name if isinstance(output_dataset_type, DatasetType) else None
+    if output_dataset_type in [DatasetType.input, DatasetType.update] or output_file_stem is None:
+        raise ValueError(
+            f"Invalid output dataset type: {output_dataset_type}. Expected one of: sym_output, asym_output, sc_output."
+        )
+
+    output_path.mkdir(parents=True, exist_ok=True)
+    json_serialize_to_file(file_path=output_path / "input.json", data=input_data, dataset_type=DatasetType.input)
+    (output_path / "input.json.license").write_text(data=LICENSE_TEXT, encoding="utf-8")
+
+    if update_data is not None:
+        json_serialize_to_file(
+            file_path=output_path / "update_batch.json", data=update_data, dataset_type=DatasetType.update
+        )
+        (output_path / "update_batch.json.license").write_text(data=LICENSE_TEXT, encoding="utf-8")
+        output_file_stem += "_batch"
+    json_serialize_to_file(
+        file_path=output_path / f"{output_file_stem}.json", data=output_data, dataset_type=output_dataset_type
+    )
+    (output_path / f"{output_file_stem}.json.license").write_text(data=LICENSE_TEXT, encoding="utf-8")
+
+    params_json = json.dumps(params, indent=2)
+    (output_path / "params.json").write_text(data=params_json, encoding="utf-8")
+    (output_path / "params.json.license").write_text(data=LICENSE_TEXT, encoding="utf-8")

tests/unit/test_utils.py

@@ -2,16 +2,19 @@
 #
 # SPDX-License-Identifier: MPL-2.0
 
+import json
 from pathlib import Path
 from unittest.mock import MagicMock, mock_open, patch
 
 import numpy as np
 import pytest
 
-from power_grid_model import initialize_array
+from power_grid_model import DatasetType, initialize_array
 from power_grid_model._core.power_grid_meta import power_grid_meta_data
 from power_grid_model.data_types import Dataset
 from power_grid_model.utils import (
+    LICENSE_TEXT,
+    _make_test_case,
     get_component_batch_size,
     get_dataset_batch_size,
     get_dataset_scenario,
@@ -164,3 +167,54 @@ def test_msgpack_serialize(serialize_mock: MagicMock, open_mock: MagicMock):
 
 def test_self_test():
     self_test()
+
+
+@pytest.mark.parametrize(
+    ("output_dataset_type", "output_file_name", "update_data"),
+    (
+        (DatasetType.sym_output, "sym_output_batch.json", {"version": "1.0", "data": "update_data"}),
+        (DatasetType.sym_output, "sym_output.json", None),
+        (DatasetType.asym_output, "asym_output_batch.json", {"version": "1.0", "data": "update_data"}),
+        (DatasetType.asym_output, "asym_output.json", None),
+        (DatasetType.sc_output, "sc_output_batch.json", {"version": "1.0", "data": "update_data"}),
+        (DatasetType.sc_output, "sc_output.json", None),
+    ),
+)
+@patch.object(Path, "write_text", autospec=True)
+@patch("power_grid_model.utils.json_serialize_to_file")
+def test__make_test_case(
+    serialize_to_file_mock: MagicMock, write_text_mock: MagicMock, output_dataset_type, output_file_name, update_data
+):
+    input_data: Dataset = {"version": "1.0", "data": "input_data"}
+    output_data: Dataset = {"version": "1.0", "data": "output_data"}
+    output_path = Path("test_path")
+    params = {"param1": "value1", "param2": "value2"}
+
+    _make_test_case(
+        output_path=output_path,
+        input_data=input_data,
+        output_data=output_data,
+        params=params,
+        output_dataset_type=output_dataset_type,
+        update_data=update_data,
+    )
+
+    serialize_to_file_mock.assert_any_call(
+        file_path=output_path / "input.json", data=input_data, dataset_type=DatasetType.input
+    )
+    serialize_to_file_mock.assert_any_call(
+        file_path=output_path / output_file_name, data=output_data, dataset_type=output_dataset_type
+    )
+    write_text_mock.assert_any_call(output_path / "params.json", data=json.dumps(params, indent=2), encoding="utf-8")
+    for file_name in ["input.json.license", f"{output_file_name}.license", "params.json.license"]:
+        write_text_mock.assert_any_call(output_path / file_name, data=LICENSE_TEXT, encoding="utf-8")
+    if update_data is not None:
+        write_text_mock.assert_any_call(output_path / "update_batch.json.license", data=LICENSE_TEXT, encoding="utf-8")
+        serialize_to_file_mock.assert_any_call(
+            file_path=output_path / "update_batch.json", data=update_data, dataset_type=DatasetType.update
+        )
+        assert write_text_mock.call_count == 5
+        assert serialize_to_file_mock.call_count == 3
+    else:
+        assert write_text_mock.call_count == 4
+        assert serialize_to_file_mock.call_count == 2