87d8873bd875820a4be6ffdef48aff86b913bf00,studio/tests/local_worker_test.py,LocalWorkerTest,test_save_get_model,#LocalWorkerTest#,125

Before Change

                     "keras is required for this test")
    def test_save_get_model(self):
        experiment_name = "test_save_get_model" + str(uuid.uuid4())
        db = stubtest_worker(
            self,
            experiment_name=experiment_name,
            runner_args=[],
            config_name="test_config.yaml",
            test_script="save_model.py",
            script_args=[],
            expected_output="",
            delete_when_done=False,
            test_output=False
        )

        experiment = db.get_experiment(experiment_name)
        saved_model = experiment.get_model(db)

After Change

                     "keras is required for this test")
    def test_save_get_model(self):
        experiment_name = "test_save_get_model" + str(uuid.uuid4())
        with stubtest_worker(
            self,
            experiment_name=experiment_name,
            runner_args=[],
            config_name="test_config.yaml",
            test_script="save_model.py",
            script_args=[],
            expected_output="",
            delete_when_done=False,
            test_output=False
        ) as db:

            experiment = db.get_experiment(experiment_name)
            saved_model = experiment.get_model(db)

            v = np.random.rand(1, 2)
            prediction = saved_model.predict(v)
            expected = v * 2

            self.assertTrue(np.isclose(prediction, expected).all())

            db.delete_experiment(experiment)

    @timeout(120)
    def test_stop_experiment(self):
        my_path = os.path.dirname(os.path.realpath(__file__))

Instances