Consolidate ax/api and ax/benchmark tests

Summary: Part of a 19-diff stack to consolidate repetitive tests across Ax and PTS using `subTest`. Consolidate 11 test files in ax/api/ and ax/benchmark/ — adds subTest to benchmark result, benchmark metric, and API utils tests.

Differential Revision: D95606935

Author: Eric Onofrey

ax/api/utils/instantiation/tests/test_from_config.py

@@ -349,21 +349,16 @@ def test_experiment_from_config(self) -> None:
         )
 
     def test_parameter_type_converter(self) -> None:
-        self.assertEqual(
-            _parameter_type_converter(parameter_type="bool"),
-            CoreParameterType.BOOL,
-        )
-        self.assertEqual(
-            _parameter_type_converter(parameter_type="int"),
-            CoreParameterType.INT,
-        )
-        self.assertEqual(
-            _parameter_type_converter(parameter_type="float"),
-            CoreParameterType.FLOAT,
-        )
-        self.assertEqual(
-            _parameter_type_converter(parameter_type="str"),
-            CoreParameterType.STRING,
-        )
+        for type_str, expected in [
+            ("bool", CoreParameterType.BOOL),
+            ("int", CoreParameterType.INT),
+            ("float", CoreParameterType.FLOAT),
+            ("str", CoreParameterType.STRING),
+        ]:
+            with self.subTest(parameter_type=type_str):
+                self.assertEqual(
+                    _parameter_type_converter(parameter_type=type_str),
+                    expected,
+                )
         with self.assertRaisesRegex(UserInputError, "Unsupported parameter type"):
             _parameter_type_converter(parameter_type="bad")

ax/benchmark/tests/problems/surrogate/hss/test_cifar10_surrogate.py

@@ -68,7 +68,8 @@ def test_cifar10_surrogate(self) -> None:
         benchmark = get_cifar10_surrogate_benchmark(num_trials=1)
 
         for params, target_value in cases:
-            self.assertAlmostEqual(
-                benchmark.test_function.evaluate_true(params).item(),
-                target_value,
-            )
+            with self.subTest(params=params):
+                self.assertAlmostEqual(
+                    benchmark.test_function.evaluate_true(params).item(),
+                    target_value,
+                )

ax/benchmark/tests/problems/surrogate/hss/test_fashion_mnist_surrogate.py

@@ -68,7 +68,8 @@ def test_fashion_mnist_surrogate(self) -> None:
         benchmark = get_fashion_mnist_surrogate_benchmark(num_trials=1)
 
         for params, target_value in cases:
-            self.assertAlmostEqual(
-                benchmark.test_function.evaluate_true(params).item(),
-                target_value,
-            )
+            with self.subTest(params=params):
+                self.assertAlmostEqual(
+                    benchmark.test_function.evaluate_true(params).item(),
+                    target_value,
+                )

ax/benchmark/tests/problems/surrogate/hss/test_mnist_surrogate.py

@@ -52,10 +52,11 @@ def test_mnist_surrogate(self) -> None:
         benchmark = get_mnist_surrogate_benchmark(num_trials=1)
 
         for params, target_value in cases:
-            self.assertAlmostEqual(
-                benchmark.test_function.evaluate_true(params).item(),
-                target_value,
-            )
+            with self.subTest(params=params):
+                self.assertAlmostEqual(
+                    benchmark.test_function.evaluate_true(params).item(),
+                    target_value,
+                )
 
     def test_benchmark_creation(self) -> None:
         benchmark = get_mnist_surrogate_benchmark(num_trials=1)

ax/benchmark/tests/problems/synthetic/hss/test_jenatton.py

@@ -78,19 +78,20 @@ def test_jenatton_test_function(self) -> None:
             )
 
         for params, value in cases:
-            self.assertAlmostEqual(
-                # pyre-fixme: Incompatible parameter type [6]: In call
-                # `jenatton_test_function`, for 1st positional argument,
-                # expected `Optional[float]` but got `Union[None, bool, float,
-                # int, str]`.
-                jenatton_test_function(**params),
-                value,
-            )
-            self.assertAlmostEqual(
-                benchmark_problem.test_function.evaluate_true(params=params).item(),
-                value,
-                places=6,
-            )
+            with self.subTest(params=params, expected_value=value):
+                self.assertAlmostEqual(
+                    # pyre-fixme: Incompatible parameter type [6]: In call
+                    # `jenatton_test_function`, for 1st positional argument,
+                    # expected `Optional[float]` but got `Union[None, bool, float,
+                    # int, str]`.
+                    jenatton_test_function(**params),
+                    value,
+                )
+                self.assertAlmostEqual(
+                    benchmark_problem.test_function.evaluate_true(params=params).item(),
+                    value,
+                    places=6,
+                )
 
     def test_create_problem(self) -> None:
         problem = get_jenatton_benchmark_problem()

ax/benchmark/tests/problems/synthetic/test_from_botorch.py

@@ -74,59 +74,65 @@ def test_get_augmented_branin_problem(self) -> None:
 class TestFromBoTorch(TestCase):
     def test_single_objective_from_botorch(self) -> None:
         for botorch_test_problem in [Ackley(), ConstrainedHartmann(dim=6)]:
-            test_problem = create_problem_from_botorch(
-                test_problem_class=botorch_test_problem.__class__,
-                test_problem_kwargs={},
-                num_trials=1,
-                baseline_value=100.0,
-            )
-
-            # Test search space
-            self.assertEqual(
-                len(test_problem.search_space.parameters), botorch_test_problem.dim
-            )
-            self.assertEqual(
-                len(test_problem.search_space.parameters),
-                len(test_problem.search_space.range_parameters),
-            )
-            self.assertTrue(
-                all(
-                    test_problem.search_space.range_parameters[f"x{i}"].lower
-                    == botorch_test_problem._bounds[i][0]
-                    for i in range(botorch_test_problem.dim)
-                ),
-                "Parameters' lower bounds must all match Botorch problem's bounds.",
-            )
-            self.assertTrue(
-                all(
-                    test_problem.search_space.range_parameters[f"x{i}"].upper
-                    == botorch_test_problem._bounds[i][1]
-                    for i in range(botorch_test_problem.dim)
-                ),
-                "Parameters' upper bounds must all match Botorch problem's bounds.",
-            )
+            with self.subTest(problem=botorch_test_problem.__class__.__name__):
+                test_problem = create_problem_from_botorch(
+                    test_problem_class=botorch_test_problem.__class__,
+                    test_problem_kwargs={},
+                    num_trials=1,
+                    baseline_value=100.0,
+                )
 
-            # Test optimum
-            self.assertEqual(
-                test_problem.optimal_value, botorch_test_problem._optimal_value
-            )
-            # test optimization config
-            metric_name = test_problem.optimization_config.objective.metric.name
-            self.assertEqual(metric_name, test_problem.name)
-            self.assertTrue(test_problem.optimization_config.objective.minimize)
-            # test repr method
-            if isinstance(botorch_test_problem, Ackley):
+                # Test search space
+                self.assertEqual(
+                    len(test_problem.search_space.parameters),
+                    botorch_test_problem.dim,
+                )
                 self.assertEqual(
-                    test_problem.optimization_config.outcome_constraints, []
+                    len(test_problem.search_space.parameters),
+                    len(test_problem.search_space.range_parameters),
                 )
-            else:
-                outcome_constraint = (
-                    test_problem.optimization_config.outcome_constraints[0]
+                self.assertTrue(
+                    all(
+                        test_problem.search_space.range_parameters[f"x{i}"].lower
+                        == botorch_test_problem._bounds[i][0]
+                        for i in range(botorch_test_problem.dim)
+                    ),
+                    "Parameters' lower bounds must all match Botorch problem's bounds.",
+                )
+                self.assertTrue(
+                    all(
+                        test_problem.search_space.range_parameters[f"x{i}"].upper
+                        == botorch_test_problem._bounds[i][1]
+                        for i in range(botorch_test_problem.dim)
+                    ),
+                    "Parameters' upper bounds must all match Botorch problem's bounds.",
+                )
+
+                # Test optimum
+                self.assertEqual(
+                    test_problem.optimal_value,
+                    botorch_test_problem._optimal_value,
                 )
-                self.assertEqual(outcome_constraint.metric.name, "constraint_slack_0")
-                self.assertEqual(outcome_constraint.op, ComparisonOp.GEQ)
-                self.assertFalse(outcome_constraint.relative)
-                self.assertEqual(outcome_constraint.bound, 0.0)
+                # test optimization config
+                metric_name = test_problem.optimization_config.objective.metric.name
+                self.assertEqual(metric_name, test_problem.name)
+                self.assertTrue(test_problem.optimization_config.objective.minimize)
+                # test repr method
+                if isinstance(botorch_test_problem, Ackley):
+                    self.assertEqual(
+                        test_problem.optimization_config.outcome_constraints,
+                        [],
+                    )
+                else:
+                    outcome_constraint = (
+                        test_problem.optimization_config.outcome_constraints[0]
+                    )
+                    self.assertEqual(
+                        outcome_constraint.metric.name, "constraint_slack_0"
+                    )
+                    self.assertEqual(outcome_constraint.op, ComparisonOp.GEQ)
+                    self.assertFalse(outcome_constraint.relative)
+                    self.assertEqual(outcome_constraint.bound, 0.0)
 
     def _test_constrained_from_botorch(
         self,
@@ -259,8 +265,9 @@ def _test_moo_from_botorch(self, lower_is_better: bool) -> None:
         )
 
     def test_moo_from_botorch(self) -> None:
-        self._test_moo_from_botorch(lower_is_better=True)
-        self._test_moo_from_botorch(lower_is_better=False)
+        for lower_is_better in [True, False]:
+            with self.subTest(lower_is_better=lower_is_better):
+                self._test_moo_from_botorch(lower_is_better=lower_is_better)
 
     def test_create_problem_from_botorch_with_shifted_function(self) -> None:
         ax_problem = create_problem_from_botorch(

ax/benchmark/tests/problems/test_mixed_integer_problems.py

@@ -29,28 +29,31 @@ def test_problems(self) -> None:
             (Rosenbrock, get_discrete_rosenbrock, 10, 6),
         ):
             name = problem_cls.__name__
-            problem = constructor()
-            self.assertEqual(f"Discrete {name}", problem.name)
-            test_function = assert_is_instance(
-                problem.test_function, BoTorchTestFunction
-            )
-            botorch_problem = test_function.botorch_problem
-            self.assertIsInstance(botorch_problem, problem_cls)
-            self.assertEqual(len(problem.search_space.parameters), dim)
-            self.assertEqual(
-                sum(
-                    p.parameter_type == ParameterType.INT
-                    for p in problem.search_space.parameters.values()
-                ),
-                dim_int,
-            )
-            # Check that the underlying problem has the correct bounds.
-            if name == "Rosenbrock":
-                expected_bounds = [(-5.0, 10.0) for _ in range(dim)]
-            else:
-                expected_bounds = [(0.0, 1.0) for _ in range(dim)]
-            self.assertEqual(botorch_problem._bounds, expected_bounds)
-            self.assertGreaterEqual(problem.optimal_value, problem_cls().optimal_value)
+            with self.subTest(problem=name):
+                problem = constructor()
+                self.assertEqual(f"Discrete {name}", problem.name)
+                test_function = assert_is_instance(
+                    problem.test_function, BoTorchTestFunction
+                )
+                botorch_problem = test_function.botorch_problem
+                self.assertIsInstance(botorch_problem, problem_cls)
+                self.assertEqual(len(problem.search_space.parameters), dim)
+                self.assertEqual(
+                    sum(
+                        p.parameter_type == ParameterType.INT
+                        for p in problem.search_space.parameters.values()
+                    ),
+                    dim_int,
+                )
+                # Check that the underlying problem has the correct bounds.
+                if name == "Rosenbrock":
+                    expected_bounds = [(-5.0, 10.0) for _ in range(dim)]
+                else:
+                    expected_bounds = [(0.0, 1.0) for _ in range(dim)]
+                self.assertEqual(botorch_problem._bounds, expected_bounds)
+                self.assertGreaterEqual(
+                    problem.optimal_value, problem_cls().optimal_value
+                )
 
         # Test that they match correctly to the original problems.
         cases: list[tuple[BenchmarkProblem, dict[str, float], torch.Tensor]] = [
@@ -94,14 +97,15 @@ def test_problems(self) -> None:
         ]
 
         for problem, params, expected_arg in cases:
-            test_function = assert_is_instance(
-                problem.test_function, BoTorchTestFunction
-            )
-            with patch.object(
-                test_function.botorch_problem,
-                attribute="evaluate_true",
-                wraps=test_function.botorch_problem.evaluate_true,
-            ) as mock_call:
-                test_function.evaluate_true(params=params)
-            actual = mock_call.call_args.kwargs["X"]
-            self.assertAllClose(actual, expected_arg)
+            with self.subTest(problem=problem.name, params=params):
+                test_function = assert_is_instance(
+                    problem.test_function, BoTorchTestFunction
+                )
+                with patch.object(
+                    test_function.botorch_problem,
+                    attribute="evaluate_true",
+                    wraps=test_function.botorch_problem.evaluate_true,
+                ) as mock_call:
+                    test_function.evaluate_true(params=params)
+                actual = mock_call.call_args.kwargs["X"]
+                self.assertAllClose(actual, expected_arg)

ax/benchmark/tests/problems/test_problems.py

@@ -64,8 +64,9 @@ def test_name(self) -> None:
             for name in ["Discrete Ackley", "Discrete Hartmann", "Discrete Rosenbrock"]
         ]
         for registry_key, problem_name in expected_names:
-            problem = get_benchmark_problem(problem_key=registry_key)
-            self.assertEqual(problem.name, problem_name)
+            with self.subTest(registry_key=registry_key):
+                problem = get_benchmark_problem(problem_key=registry_key)
+                self.assertEqual(problem.name, problem_name)
 
     def test_no_duplicates(self) -> None:
         problem_names = set()

ax/benchmark/tests/test_benchmark.py

@@ -189,8 +189,9 @@ def _test_storage(self, map_data: bool) -> None:
             self.assertEqual(experiment, experiment)
 
     def test_storage(self) -> None:
-        self._test_storage(map_data=False)
-        self._test_storage(map_data=True)
+        for map_data in [False, True]:
+            with self.subTest(map_data=map_data):
+                self._test_storage(map_data=map_data)
 
     def test_replication_sobol_synthetic(self) -> None:
         method = get_sobol_benchmark_method()

ax/benchmark/tests/test_benchmark_metric.py

@@ -341,8 +341,9 @@ def _test_fetch_trial_multiple_time_steps_with_simulator(self, batch: bool) -> N
             self.assertEqual(df.reset_index(drop=True).to_dict(), expected_df.to_dict())
 
     def test_fetch_trial_multiple_time_steps_with_simulator(self) -> None:
-        self._test_fetch_trial_multiple_time_steps_with_simulator(batch=False)
-        self._test_fetch_trial_multiple_time_steps_with_simulator(batch=True)
+        for batch in [False, True]:
+            with self.subTest(batch=batch):
+                self._test_fetch_trial_multiple_time_steps_with_simulator(batch=batch)
 
     def test_sim_trial_completes_in_future_raises(self) -> None:
         simulator = BackendSimulator()