Fix unittests.

lbl-srg · Dec 18, 2017 · a3d58b7 · a3d58b7
1 parent 57bc0fd
commit a3d58b7
Showing 1 changed file with 7 additions and 11 deletions.
diff --git a/unittests/test_optimization.py b/unittests/test_optimization.py
@@ -220,17 +220,13 @@ def test_set_options_error(self):
                                                 optimization.JModelica, \
                                                 'q_flow', \
                                                 constraint_data = self.constraints.data);
-        # Get initial options
-        opt_options = opt_problem.get_optimization_options();
-        # Set new options
-        opt_options['n_e'] = 2;
-        self.assertRaises(KeyError, opt_problem.set_optimization_options(opt_options));
-        opt_options['external_data'] = 2;
-        self.assertRaises(KeyError, opt_problem.set_optimization_options(opt_options));
-        opt_options['init_traj'] = 2;
-        self.assertRaises(KeyError, opt_problem.set_optimization_options(opt_options));
-        opt_options['nominal_traj'] = 2;
-        self.assertRaises(KeyError, opt_problem.set_optimization_options(opt_options));
+        for key in ['external_data', 'init_traj', 'nominal_traj']:
+            # Get initial options
+            opt_options = opt_problem.get_optimization_options();
+            # Set new options and check KeyError raised
+            with self.assertRaises(KeyError):
+                opt_options[key] = 2;
+                opt_problem.set_optimization_options(opt_options);
 
     def test_get_statistics(self):
         '''Test the getting of optimization result statistics.