Consolidate ax/adapter tests (#5006)

Summary:

Part of a 19-diff stack to consolidate repetitive tests across Ax and PTS using `subTest`. Consolidate 8 test files in ax/adapter/ and ax/adapter/transforms/ — adds subTest parameterization to torch adapter, trial-as-task transform, logit transform, and objective-as-constraint tests.

Differential Revision: D95603401

Author: eonofrey

ax/adapter/tests/test_adapter_utils.py

@@ -313,33 +313,16 @@ def get_adapter(min: float, max: float) -> TorchAdapter:
 
     def test_arm_to_np_array(self) -> None:
         # Test extracting target point from arm with valid parameters
-
-        # Setup: create arm with target parameter values
         target_arm = Arm(parameters={"x1": 0.5, "x2": 1.5, "x3": 2.5})
-        parameters = ["x1", "x2", "x3"]
-
-        # Execute: extract target point
-        actual = arm_to_np_array(arm=target_arm, parameters=parameters)
-
-        # Assert: confirm extracted values match expected order
-        expected = np.array([0.5, 1.5, 2.5])
-        self.assertIsNotNone(actual)
-        np.testing.assert_array_equal(actual, expected)
-
-    def test_extract_arm_to_np_array_different_parameter_order(self) -> None:
-        # Test extracting target point with different parameter ordering
-
-        # Setup: create arm and specify parameters in different order
-        target_arm = Arm(parameters={"x1": 0.5, "x2": 1.5, "x3": 2.5})
-        parameters = ["x3", "x1", "x2"]
-
-        # Execute: extract target point
-        actual = arm_to_np_array(arm=target_arm, parameters=parameters)
-
-        # Assert: confirm values are extracted in specified parameter order
-        expected = np.array([2.5, 0.5, 1.5])
-        self.assertIsNotNone(actual)
-        np.testing.assert_array_equal(actual, expected)
+        cases = [
+            (["x1", "x2", "x3"], np.array([0.5, 1.5, 2.5])),
+            (["x3", "x1", "x2"], np.array([2.5, 0.5, 1.5])),
+        ]
+        for parameters, expected in cases:
+            with self.subTest(parameters=parameters):
+                actual = arm_to_np_array(arm=target_arm, parameters=parameters)
+                self.assertIsNotNone(actual)
+                np.testing.assert_array_equal(actual, expected)
 
     def test_arm_to_np_array_none(self) -> None:
         # Test that None is returned when target_arm is None

ax/adapter/tests/test_base_adapter.py

@@ -193,11 +193,10 @@ def _test_init_with_data(self, multi_objective: bool) -> None:
             search_space=search_space, experiment_data=experiment_data
         )
 
-    def test_init_with_data_single_objective(self) -> None:
-        self._test_init_with_data(multi_objective=False)
-
-    def test_init_with_data_multi_objective(self) -> None:
-        self._test_init_with_data(multi_objective=True)
+    def test_init_with_data(self) -> None:
+        for multi_objective in (False, True):
+            with self.subTest(multi_objective=multi_objective):
+                self._test_init_with_data(multi_objective=multi_objective)
 
     def test_fit_tracking_metrics(self) -> None:
         # Test error when fit_tracking_metrics is False and optimization

ax/adapter/tests/test_hierarchical_search_space.py

@@ -220,20 +220,16 @@ def _base_test_predict_and_cv(
         cv_res = cross_validate(adapter=mbm)
         self.assertEqual(len(cv_res), len(experiment.trials))
 
-    def test_with_non_hierarchical_hss(self) -> None:
-        experiment = self._test_gen_base(
-            hss=self.non_hierarchical_hss, expected_num_candidate_params=[3]
-        )
-        self._base_test_predict_and_cv(experiment=experiment)
-
-    def test_with_simple_hss(self) -> None:
-        experiment = self._test_gen_base(
-            hss=self.simple_hss, expected_num_candidate_params=[2]
-        )
-        self._base_test_predict_and_cv(experiment=experiment)
-
-    def test_with_complex_hss(self) -> None:
-        experiment = self._test_gen_base(
-            hss=self.complex_hss, expected_num_candidate_params=[2, 4, 5]
-        )
-        self._base_test_predict_and_cv(experiment=experiment)
+    def test_with_hss_variants(self) -> None:
+        cases = [
+            ("non_hierarchical", self.non_hierarchical_hss, [3]),
+            ("simple", self.simple_hss, [2]),
+            ("complex", self.complex_hss, [2, 4, 5]),
+        ]
+        for label, hss, expected_num_candidate_params in cases:
+            with self.subTest(hss_variant=label):
+                experiment = self._test_gen_base(
+                    hss=hss,
+                    expected_num_candidate_params=expected_num_candidate_params,
+                )
+                self._base_test_predict_and_cv(experiment=experiment)

ax/adapter/tests/test_torch_adapter.py

@@ -1251,90 +1251,35 @@ def test_pairwise_preference_generator(self) -> None:
                 X=X.expand(2, *X.shape), Y=comp_pair_Y.expand(2, *comp_pair_Y.shape)
             )
 
-    def test_get_transformed_model_gen_args_with_target_point(self) -> None:
-        # Test that _get_transformed_model_gen_args correctly processes target_point
-
-        # Setup: create adapter with target arm in optimization config
-        experiment = get_branin_experiment(with_completed_trial=True)
-        pruning_target_parameterization = Arm(parameters={"x1": -5.0, "x2": 15.0})
-        optimization_config = none_throws(
-            experiment.optimization_config
-        ).clone_with_args(
-            pruning_target_parameterization=pruning_target_parameterization
-        )
-
-        adapter = TorchAdapter(
-            generator=TorchGenerator(),
-            experiment=experiment,
-            transforms=Cont_X_trans,
-        )
-
-        # Execute: call _get_transformed_gen_args then _get_transformed_model_gen_args
-        base_gen_args = adapter._get_transformed_gen_args(
-            search_space=experiment.search_space,
-            optimization_config=optimization_config,
-            pending_observations={},
-        )
-
-        search_space_digest, torch_opt_config = adapter._get_transformed_model_gen_args(
-            search_space=base_gen_args.search_space,
-            pending_observations=base_gen_args.pending_observations,
-            fixed_features=base_gen_args.fixed_features,
-            optimization_config=base_gen_args.optimization_config,
-        )
-
-        # Assert: confirm pruning_target_point is correctly extracted and transformed
-        self.assertIsNotNone(torch_opt_config.pruning_target_point)
-        expected_target = torch.tensor([0.0, 1.0], dtype=torch.double)
-        torch.testing.assert_close(
-            torch_opt_config.pruning_target_point, expected_target
+    def _test_get_transformed_model_gen_args_target_point(
+        self,
+        with_status_quo: bool,
+        pruning_target_params: dict[str, float] | None,
+        expected_target: torch.Tensor | None,
+    ) -> None:
+        experiment = get_branin_experiment(
+            with_completed_trial=True,
+            with_status_quo=with_status_quo,
         )
 
-    def test_get_transformed_model_gen_args_no_target_point(self) -> None:
-        # Test that _get_transformed_model_gen_args handles
-        # pruning_target_parameterization=None correctly
+        opt_config = none_throws(experiment.optimization_config)
+        if pruning_target_params is not None:
+            pruning_target = Arm(parameters=pruning_target_params)
+            opt_config = opt_config.clone_with_args(
+                pruning_target_parameterization=pruning_target
+            )
+        elif with_status_quo:
+            opt_config = opt_config.clone()
 
-        # Setup: create adapter without target arm (default case)
-        experiment = get_branin_experiment(with_completed_trial=True)
         adapter = TorchAdapter(
             generator=TorchGenerator(),
             experiment=experiment,
             transforms=Cont_X_trans,
         )
 
-        # Execute: call _get_transformed_gen_args then _get_transformed_model_gen_args
         base_gen_args = adapter._get_transformed_gen_args(
             search_space=experiment.search_space,
-            optimization_config=none_throws(experiment.optimization_config),
-            pending_observations={},
-        )
-
-        search_space_digest, torch_opt_config = adapter._get_transformed_model_gen_args(
-            search_space=base_gen_args.search_space,
-            pending_observations=base_gen_args.pending_observations,
-            fixed_features=base_gen_args.fixed_features,
-            optimization_config=base_gen_args.optimization_config,
-        )
-
-        # Assert: confirm target_point is None when no pruning_target_parameterization
-        #  is provided
-        self.assertIsNone(torch_opt_config.pruning_target_point)
-
-    def test_get_transformed_model_gen_args_with_sq_as_target(self) -> None:
-        # Test that _get_transformed_model_gen_args correctly processes the status quo
-        # as the target point
-        experiment = get_branin_experiment(
-            with_completed_trial=True, with_status_quo=True
-        )
-
-        adapter = TorchAdapter(
-            generator=TorchGenerator(), experiment=experiment, transforms=Cont_X_trans
-        )
-        oc = none_throws(experiment.optimization_config).clone()
-        # Execute: call _get_transformed_gen_args then _get_transformed_model_gen_args
-        base_gen_args = adapter._get_transformed_gen_args(
-            search_space=experiment.search_space,
-            optimization_config=oc,
+            optimization_config=opt_config,
             pending_observations={},
         )
 
@@ -1345,12 +1290,38 @@ def test_get_transformed_model_gen_args_with_sq_as_target(self) -> None:
             optimization_config=base_gen_args.optimization_config,
         )
 
-        # Assert: confirm pruning_target_point is correctly extracted and transformed
-        self.assertIsNotNone(torch_opt_config.pruning_target_point)
-        expected_target = torch.tensor([1 / 3.0, 0.0], dtype=torch.double)
-        torch.testing.assert_close(
-            torch_opt_config.pruning_target_point, expected_target
-        )
+        if expected_target is None:
+            self.assertIsNone(torch_opt_config.pruning_target_point)
+        else:
+            self.assertIsNotNone(torch_opt_config.pruning_target_point)
+            torch.testing.assert_close(
+                torch_opt_config.pruning_target_point,
+                expected_target,
+            )
+
+    def test_get_transformed_model_gen_args_target_point(self) -> None:
+        # Test _get_transformed_model_gen_args with various target point scenarios
+        for label, with_status_quo, pruning_target_params, expected_target in [
+            (
+                "with_target_point",
+                False,
+                {"x1": -5.0, "x2": 15.0},
+                torch.tensor([0.0, 1.0], dtype=torch.double),
+            ),
+            ("no_target_point", False, None, None),
+            (
+                "sq_as_target",
+                True,
+                None,
+                torch.tensor([1 / 3.0, 0.0], dtype=torch.double),
+            ),
+        ]:
+            with self.subTest(scenario=label):
+                self._test_get_transformed_model_gen_args_target_point(
+                    with_status_quo=with_status_quo,
+                    pruning_target_params=pruning_target_params,
+                    expected_target=expected_target,
+                )
 
     @mock_botorch_optimize
     def test_moo_with_derived_parameter(self) -> None:

ax/adapter/transforms/tests/test_logit_transform.py

@@ -95,15 +95,11 @@ def test_InvalidSettings(self) -> None:
         self.assertEqual("x can't use both log and logit.", str(cm.exception))
 
         str_exc = "x logit requires lower > 0 and upper < 1"
-        with self.assertRaises(UserInputError) as cm:
-            self._create_logit_parameter(lower=0.0, upper=0.5)
-        self.assertEqual(str_exc, str(cm.exception))
-        with self.assertRaises(UserInputError) as cm:
-            self._create_logit_parameter(lower=0.3, upper=1.0)
-        self.assertEqual(str_exc, str(cm.exception))
-        with self.assertRaises(UserInputError) as cm:
-            self._create_logit_parameter(lower=0.5, upper=10.0)
-        self.assertEqual(str_exc, str(cm.exception))
+        for lower, upper in [(0.0, 0.5), (0.3, 1.0), (0.5, 10.0)]:
+            with self.subTest(lower=lower, upper=upper):
+                with self.assertRaises(UserInputError) as cm:
+                    self._create_logit_parameter(lower=lower, upper=upper)
+                self.assertEqual(str_exc, str(cm.exception))
 
     def test_TransformSearchSpace(self) -> None:
         ss2 = deepcopy(self.search_space)

ax/adapter/transforms/tests/test_map_key_to_float_transform.py

@@ -249,11 +249,10 @@ def _test_early_stopping(self, complete_with_progression: bool) -> None:
         # Check that cross validation works.
         cross_validate(adapter=adapter)
 
-    def test_no_early_stopping_with_progression(self) -> None:
-        self._test_no_early_stopping(with_progression=True)
-
-    def test_no_early_stopping_no_progression(self) -> None:
-        self._test_no_early_stopping(with_progression=False)
+    def test_no_early_stopping(self) -> None:
+        for with_progression in (True, False):
+            with self.subTest(with_progression=with_progression):
+                self._test_no_early_stopping(with_progression=with_progression)
 
     def test_early_stopping_with_final_progression(self) -> None:
         self._test_early_stopping(complete_with_progression=True)

ax/adapter/transforms/tests/test_objective_as_constraint.py

@@ -334,37 +334,25 @@ def test_relative_constraint_feasibility_check(self) -> None:
 
     def test_leq_constraint_feasibility(self) -> None:
         """Test feasibility checking with LEQ constraints."""
-        # m2 <= 0.3 constraint. Both observations have m2 > 0.3, so infeasible.
-        _, adapter, experiment_data = self._make_experiment_adapter_and_data(
-            observations=[[1.0, 0.5], [2.0, 5.0]],
-            constraint_bound=0.3,
-            constraint_op=ComparisonOp.LEQ,
-        )
-
-        t = ObjectiveAsConstraint(
-            search_space=adapter._experiment.search_space,
-            experiment_data=experiment_data,
-            adapter=adapter,
-        )
-
-        self.assertTrue(t._should_add_constraint)
-
-    def test_leq_constraint_feasible(self) -> None:
-        """Test that LEQ constraints with feasible points are correctly detected."""
-        # m2 <= 10.0 constraint. Both observations have m2 <= 10.0, so feasible.
-        _, adapter, experiment_data = self._make_experiment_adapter_and_data(
-            observations=[[1.0, 0.5], [2.0, 5.0]],
-            constraint_bound=10.0,
-            constraint_op=ComparisonOp.LEQ,
-        )
+        cases = [
+            (0.3, True, "infeasible"),
+            (10.0, False, "feasible"),
+        ]
+        for bound, expected_should_add, label in cases:
+            with self.subTest(bound=bound, scenario=label):
+                _, adapter, experiment_data = self._make_experiment_adapter_and_data(
+                    observations=[[1.0, 0.5], [2.0, 5.0]],
+                    constraint_bound=bound,
+                    constraint_op=ComparisonOp.LEQ,
+                )
 
-        t = ObjectiveAsConstraint(
-            search_space=adapter._experiment.search_space,
-            experiment_data=experiment_data,
-            adapter=adapter,
-        )
+                t = ObjectiveAsConstraint(
+                    search_space=adapter._experiment.search_space,
+                    experiment_data=experiment_data,
+                    adapter=adapter,
+                )
 
-        self.assertFalse(t._should_add_constraint)
+                self.assertEqual(t._should_add_constraint, expected_should_add)
 
     def test_no_op_for_experiment_data(self) -> None:
         """Test that transform_experiment_data is a no-op."""