diff --git a/tests/test_step.py b/tests/test_step.py index ea38a0b4..dade4b5e 100644 --- a/tests/test_step.py +++ b/tests/test_step.py @@ -509,6 +509,15 @@ def pop(self, index=-1): self._models.pop(index) +class SimpleContainerWithSave(SimpleContainer): + + def save(self, path, dir_path=None, *args, **kwargs): + for model in self._models[1:]: + # skip the first model to test that the save method is called + # rather than just looping over all models like in the without-save case + model.save(path, dir_path, *args, **kwargs) + + def test_save_container(tmp_cwd, model_list): """ensure list-like save still works for non-list sequence""" container = SimpleContainer(model_list) @@ -518,6 +527,16 @@ def test_save_container(tmp_cwd, model_list): assert (tmp_cwd / f"test{i}-saved.txt").exists() +def test_save_container_with_save_method(tmp_cwd, model_list): + """ensure list-like save still works for non-list sequence""" + container = SimpleContainerWithSave(model_list) + step = StepWithModel() + step.run(container) + assert not (tmp_cwd / "test0-saved.txt").exists() + assert (tmp_cwd / "test1-saved.txt").exists() + assert (tmp_cwd / "test2-saved.txt").exists() + + def test_save_tuple_with_nested_list(tmp_cwd, model_list): """ in rare cases, multiple outputs are returned from step as tuple.