Extend tests for build_spec functionality

Author: constantinpape

tests/build_spec/test_add_weights.py

@@ -1,4 +1,5 @@
 from bioimageio.core import export_resource_package, load_raw_resource_description, load_resource_description
+from bioimageio.core.resource_tests import test_model as _test_model
 
 
 def _test_add_weights(model, tmp_path, base_weights, added_weights, **kwargs):
@@ -28,6 +29,10 @@ def _test_add_weights(model, tmp_path, base_weights, added_weights, **kwargs):
     for weight in new_rdf.weights.values():
         assert weight.source.exists()
 
+    test_res = _test_model(out_path, added_weights)
+    test_res = _test_model(out_path)
+    assert test_res["error"] is None
+
 
 def test_add_torchscript(unet2d_nuclei_broad_model, tmp_path):
     _test_add_weights(unet2d_nuclei_broad_model, tmp_path, "pytorch_state_dict", "torchscript")

tests/build_spec/test_build_spec.py

@@ -3,6 +3,7 @@
 from bioimageio.core import load_raw_resource_description, load_resource_description
 from bioimageio.core.resource_io import nodes
 from bioimageio.core.resource_io.utils import resolve_source
+from bioimageio.core.resource_tests import test_model as _test_model
 from marshmallow import missing
 
 
@@ -81,17 +82,19 @@ def _test_build_spec(
         output_path=out_path,
         add_deepimagej_config=add_deepimagej_config,
         maintainers=[{"github_user": "jane_doe"}],
+        input_names=[inp.name for inp in model_spec.inputs],
+        output_names=[out.name for out in model_spec.outputs],
     )
     if architecture is not None:
         kwargs["architecture"] = architecture
     if model_kwargs is not None:
-        kwargs["kwargs"] = model_kwargs
+        kwargs["model_kwargs"] = model_kwargs
     if tensorflow_version is not None:
         kwargs["tensorflow_version"] = tensorflow_version
     if opset_version is not None:
         kwargs["opset_version"] = opset_version
     if use_implicit_output_shape:
-        kwargs["input_name"] = ["input"]
+        kwargs["input_names"] = ["input"]
         kwargs["output_reference"] = ["input"]
         kwargs["output_scale"] = [[1.0, 1.0, 1.0, 1.0]]
         kwargs["output_offset"] = [[0.0, 0.0, 0.0, 0.0]]
@@ -134,6 +137,10 @@ def _test_build_spec(
             assert attachments.files is not missing
             assert n_processing == len(attachments.files)
 
+    # test inference for the model to ensure that the weights were written correctly
+    test_res = _test_model(out_path)
+    assert test_res["error"] is None
+
 
 def test_build_spec_pytorch(any_torch_model, tmp_path):
     _test_build_spec(any_torch_model, tmp_path / "model.zip", "pytorch_state_dict")