diff --git a/test/test_modules/test_model_setup.py b/test/test_modules/test_model_setup.py index b8652c33852b56b5eaee369a4662b68b95316291..541f8d17c8084a34a75b1907c5ff89b08bea9e03 100644 --- a/test/test_modules/test_model_setup.py +++ b/test/test_modules/test_model_setup.py @@ -22,6 +22,7 @@ class TestModelSetup: obj.data_store.set("model_class", MyLittleModel) obj.data_store.set("lr_decay", "dummy_str", "general.model") obj.data_store.set("hist", "dummy_str", "general.model") + obj.data_store.set("epochs", 2) obj.model_name = "%s.h5" yield obj RunEnvironment().__del__() @@ -49,8 +50,7 @@ class TestModelSetup: @pytest.fixture def setup_with_model(self, setup): setup.model = AbstractModelClass() - setup.model.epochs = 2 - setup.model.batch_size = int(256) + setup.model.test_param = "42" yield setup RunEnvironment().__del__() @@ -80,15 +80,15 @@ class TestModelSetup: setup_with_model.scope = "model_test" with pytest.raises(EmptyScope): self.current_scope_as_set(setup_with_model) # will fail because scope is not created - setup_with_model.get_model_settings() # this saves now the parameters epochs and batch_size into scope - assert {"epochs", "batch_size"} <= self.current_scope_as_set(setup_with_model) + setup_with_model.get_model_settings() # this saves now the parameter test_param into scope + assert {"test_param", "model_name"} <= self.current_scope_as_set(setup_with_model) def test_build_model(self, setup_with_gen): assert setup_with_gen.model is None setup_with_gen.build_model() assert isinstance(setup_with_gen.model, AbstractModelClass) expected = {"window_history_size", "window_lead_time", "channels", "dropout_rate", "regularizer", "initial_lr", - "optimizer", "epochs", "batch_size", "activation"} + "optimizer", "batch_size", "activation"} assert expected <= self.current_scope_as_set(setup_with_gen) def test_set_channels(self, setup_with_gen_tiny):