Consolidate ax/core tests (#5009)

Summary:
Pull Request resolved: #5009

Part of a 19-diff stack to consolidate repetitive tests across Ax using `subTest`. Consolidate 12 test files in ax/core/ — adds subTest to experiment, metric, parameter constraint, runner, trial, and types tests.

Differential Revision: D95604288

Author: eonofrey

ax/core/tests/test_batch_trial.py

@@ -530,34 +530,39 @@ def test_Repr(self) -> None:
             "BatchTrial(experiment_name='test', index=0, status=TrialStatus.CANDIDATE)",
         )
 
-    def test_TTL_status_property_check(self) -> None:
-        """Verify that TTL is checked on execution of the `status` property."""
-        candidate_trial = self.experiment.new_batch_trial(ttl_seconds=1)
-        self.assertTrue(candidate_trial.status.is_candidate)
-        sleep(1)  # Wait 1 second for trial TTL to elapse.
-        # Candidate should become stale on expiration of TTL.
-        self.assertTrue(candidate_trial.status.is_stale)
-        self.assertIn(1, self.experiment.trial_indices_by_status[TrialStatus.STALE])
-
-    def test_TTL_trial_indices_by_status_check(self) -> None:
-        """Verify that TTL is checked on `experiment.trial_indices_by_status`."""
-        candidate_trial = self.experiment.new_batch_trial(ttl_seconds=1)
-        self.assertTrue(candidate_trial.status.is_candidate)
-        sleep(1)  # Wait 1 second for trial TTL to elapse.
-        self.assertIn(1, self.experiment.trial_indices_by_status[TrialStatus.STALE])
-        self.assertTrue(candidate_trial.status.is_stale)
-
-    def test_TTL_experiment_trials_check(self) -> None:
-        """Verify that TTL is checked on `experiment.trials`."""
-        candidate_trial = self.experiment.new_batch_trial(ttl_seconds=1)
-        self.assertTrue(candidate_trial.status.is_candidate)
-        self.assertIn(1, self.experiment.trial_indices_by_status[TrialStatus.CANDIDATE])
-        sleep(1)  # Wait 1 second for trial TTL to elapse.
-        self.experiment.trials
-        # Check `_status`, not `status`, to ensure it's within `trials` that the status
-        # was actually changed, not in `status`.
-        self.assertEqual(candidate_trial._status, TrialStatus.STALE)
-        self.assertIn(1, self.experiment.trial_indices_by_status[TrialStatus.STALE])
+    def test_TTL_expiry_across_access_paths(self) -> None:
+        """Verify that TTL is checked via different access paths."""
+
+        def check_via_experiment_trials(trial: BatchTrial) -> bool:
+            # Access experiment.trials first (this is the side-effect that
+            # triggers the TTL status update), then separately verify
+            # that the private _status field was updated.
+            _ = self.experiment.trials
+            return trial._status == TrialStatus.STALE
+
+        # Three different code paths that should all trigger TTL expiry:
+        # 1. Accessing trial.status directly
+        # 2. Querying experiment.trial_indices_by_status
+        # 3. Accessing experiment.trials (updates _status as side effect)
+        ttl_access_paths = {
+            "status_property": lambda trial: trial.status.is_stale,
+            "trial_indices_by_status": lambda trial: (
+                1 in self.experiment.trial_indices_by_status[TrialStatus.STALE]
+            ),
+            "experiment_trials": check_via_experiment_trials,
+        }
+        for access_path, check_stale in ttl_access_paths.items():
+            with self.subTest(access_path=access_path):
+                # Reset by creating a fresh experiment for each sub-test
+                self.setUp()
+                candidate_trial = self.experiment.new_batch_trial(ttl_seconds=1)
+                self.assertTrue(candidate_trial.status.is_candidate)
+                sleep(1)  # Wait 1 second for trial TTL to elapse.
+                self.assertTrue(check_stale(candidate_trial))
+                self.assertIn(
+                    1,
+                    self.experiment.trial_indices_by_status[TrialStatus.STALE],
+                )
 
     def test_get_candidate_metadata_from_all_generator_runs(self) -> None:
         self.assertEqual(self.batch.generation_method_str, MANUAL_GENERATION_METHOD_STR)

ax/core/tests/test_data.py

@@ -793,47 +793,44 @@ def setUp(self) -> None:
         )
 
     def test_relativize_data(self) -> None:
-        data = Data(df=self.df)
-        expected_relativized_data = Data(df=self.expected_relativized_df)
+        # Test relativization both with valid SEM values and with NaN SEMs
+        for sem_label, modify_sem in [
+            ("with_sem", False),
+            ("no_sem", True),
+        ]:
+            with self.subTest(sem=sem_label):
+                df = self.df.copy()
+                expected_relativized_df = self.expected_relativized_df.copy()
+                expected_relativized_df_with_sq = (
+                    self.expected_relativized_df_with_sq.copy()
+                )
 
-        expected_relativized_data_with_sq = Data(
-            df=self.expected_relativized_df_with_sq
-        )
+                if modify_sem:
+                    df["sem"] = np.nan
+                    expected_relativized_df["sem"] = np.nan
+                    expected_relativized_df_with_sq.loc[
+                        expected_relativized_df_with_sq["arm_name"] != "status_quo",
+                        "sem",
+                    ] = np.nan
 
-        actual_relativized_data = data.relativize()
-        self.assertEqual(expected_relativized_data, actual_relativized_data)
+                data = Data(df=df)
+                expected_relativized_data = Data(df=expected_relativized_df)
+                expected_relativized_data_with_sq = Data(
+                    df=expected_relativized_df_with_sq
+                )
 
-        actual_relativized_data_with_sq = data.relativize(include_sq=True)
-        self.assertEqual(
-            expected_relativized_data_with_sq, actual_relativized_data_with_sq
-        )
+                actual_relativized_data = data.relativize()
+                self.assertEqual(expected_relativized_data, actual_relativized_data)
+
+                actual_relativized_data_with_sq = data.relativize(include_sq=True)
+                self.assertEqual(
+                    expected_relativized_data_with_sq,
+                    actual_relativized_data_with_sq,
+                )
 
         with self.subTest("step column not supported"):
             data = Data(df=self.df.assign(step=0))
             with self.assertRaisesRegex(
                 NotImplementedError, "Relativization is not supported"
             ):
                 data.relativize()
-
-    def test_relativize_data_no_sem(self) -> None:
-        df = self.df.copy()
-        df["sem"] = np.nan
-        data = Data(df=df)
-
-        expected_relativized_df = self.expected_relativized_df.copy()
-        expected_relativized_df["sem"] = np.nan
-        expected_relativized_data = Data(df=expected_relativized_df)
-
-        expected_relativized_df_with_sq = self.expected_relativized_df_with_sq.copy()
-        expected_relativized_df_with_sq.loc[
-            expected_relativized_df_with_sq["arm_name"] != "status_quo", "sem"
-        ] = np.nan
-        expected_relativized_data_with_sq = Data(df=expected_relativized_df_with_sq)
-
-        actual_relativized_data = data.relativize()
-        self.assertEqual(expected_relativized_data, actual_relativized_data)
-
-        actual_relativized_data_with_sq = data.relativize(include_sq=True)
-        self.assertEqual(
-            expected_relativized_data_with_sq, actual_relativized_data_with_sq
-        )

ax/core/tests/test_experiment.py

@@ -1102,79 +1102,74 @@ def test_immutable_search_space_and_opt_config(self) -> None:
         )
         self.assertTrue(immutable_exp_2.immutable_search_space_and_opt_config)
 
-    def test_attach_batch_trial_no_arm_names(self) -> None:
-        num_trials = len(self.experiment.trials)
-
-        _, trial_index = self.experiment.attach_trial(
-            parameterizations=[
-                {"w": 5.3, "x": 5, "y": "baz", "z": True, "d": 11.6},
-                {"w": 5.2, "x": 5, "y": "foo", "z": True, "d": 11.4},
-                {"w": 5.1, "x": 5, "y": "bar", "z": True, "d": 11.2},
-            ],
-            ttl_seconds=3600,
-            run_metadata={"test_metadata_field": 1},
-        )
-
-        self.assertEqual(len(self.experiment.trials), num_trials + 1)
-        self.assertEqual(
-            len(set(self.experiment.trials[trial_index].arms_by_name) - {"status_quo"}),
-            3,
-        )
-        self.assertEqual(type(self.experiment.trials[trial_index]), BatchTrial)
-
-    def test_attach_batch_trial_with_arm_names(self) -> None:
-        num_trials = len(self.experiment.trials)
-
-        _, trial_index = self.experiment.attach_trial(
-            parameterizations=[
-                {"w": 5.3, "x": 5, "y": "baz", "z": True, "d": 11.6},
-                {"w": 5.2, "x": 5, "y": "foo", "z": True, "d": 11.4},
-                {"w": 5.1, "x": 5, "y": "bar", "z": True, "d": 11.2},
-            ],
-            arm_names=["arm1", "arm2", "arm3"],
-            ttl_seconds=3600,
-            run_metadata={"test_metadata_field": 1},
-        )
-
-        self.assertEqual(len(self.experiment.trials), num_trials + 1)
-        self.assertEqual(
-            len(set(self.experiment.trials[trial_index].arms_by_name) - {"status_quo"}),
-            3,
-        )
-        self.assertEqual(type(self.experiment.trials[trial_index]), BatchTrial)
-        self.assertEqual(
-            {"arm1", "arm2", "arm3"},
-            set(self.experiment.trials[trial_index].arms_by_name) - {"status_quo"},
-        )
-
-    def test_attach_single_arm_trial_no_arm_name(self) -> None:
+    def _test_attach_trial(
+        self,
+        parameterizations: list[dict[str, object]],
+        arm_names: list[str] | None,
+        expected_type: type[BaseTrial],
+        expected_names: set[str] | str | None,
+    ) -> None:
+        self.setUp()
         num_trials = len(self.experiment.trials)
-
-        _, trial_index = self.experiment.attach_trial(
-            parameterizations=[{"w": 5.3, "x": 5, "y": "baz", "z": True, "d": 11.6}],
-            ttl_seconds=3600,
-            run_metadata={"test_metadata_field": 1},
-        )
+        if arm_names is not None:
+            _, trial_index = self.experiment.attach_trial(
+                parameterizations=parameterizations,
+                arm_names=arm_names,
+                ttl_seconds=3600,
+                run_metadata={"test_metadata_field": 1},
+            )
+        else:
+            _, trial_index = self.experiment.attach_trial(
+                parameterizations=parameterizations,
+                ttl_seconds=3600,
+                run_metadata={"test_metadata_field": 1},
+            )
 
         self.assertEqual(len(self.experiment.trials), num_trials + 1)
-        self.assertEqual(type(self.experiment.trials[trial_index]), Trial)
-
-    def test_attach_single_arm_trial_with_arm_name(self) -> None:
-        num_trials = len(self.experiment.trials)
-
-        _, trial_index = self.experiment.attach_trial(
-            parameterizations=[{"w": 5.3, "x": 5, "y": "baz", "z": True, "d": 11.6}],
-            arm_names=["arm1"],
-            ttl_seconds=3600,
-            run_metadata={"test_metadata_field": 1},
-        )
+        self.assertEqual(type(self.experiment.trials[trial_index]), expected_type)
+        if isinstance(expected_names, set):
+            self.assertEqual(
+                expected_names,
+                set(self.experiment.trials[trial_index].arms_by_name) - {"status_quo"},
+            )
+        elif isinstance(expected_names, str):
+            self.assertEqual(
+                expected_names,
+                self.experiment.trials[trial_index].arm.name,
+            )
 
-        self.assertEqual(len(self.experiment.trials), num_trials + 1)
-        self.assertEqual(type(self.experiment.trials[trial_index]), Trial)
-        self.assertEqual(
-            "arm1",
-            self.experiment.trials[trial_index].arm.name,
-        )
+    def test_attach_trial(self) -> None:
+        batch_params: list[dict[str, object]] = [
+            {"w": 5.3, "x": 5, "y": "baz", "z": True, "d": 11.6},
+            {"w": 5.2, "x": 5, "y": "foo", "z": True, "d": 11.4},
+            {"w": 5.1, "x": 5, "y": "bar", "z": True, "d": 11.2},
+        ]
+        single_params: list[dict[str, object]] = [
+            {"w": 5.3, "x": 5, "y": "baz", "z": True, "d": 11.6},
+        ]
+        for label, params, arm_names, expected_type, expected_names in [
+            # Batch trial without explicit arm names -> auto-generated names
+            ("batch_no_arm_names", batch_params, None, BatchTrial, None),
+            # Batch trial with explicit arm names -> names preserved
+            (
+                "batch_with_arm_names",
+                batch_params,
+                ["arm1", "arm2", "arm3"],
+                BatchTrial,
+                {"arm1", "arm2", "arm3"},
+            ),
+            # Single-arm trial without arm name -> creates Trial (not BatchTrial)
+            ("single_no_arm_name", single_params, None, Trial, None),
+            # Single-arm trial with explicit arm name -> name preserved
+            ("single_with_arm_name", single_params, ["arm1"], Trial, "arm1"),
+        ]:
+            with self.subTest(label=label):
+                self._test_attach_trial(
+                    parameterizations=params,
+                    arm_names=arm_names,
+                    expected_type=expected_type,
+                    expected_names=expected_names,
+                )
 
     def test_fetch_as_class(self) -> None:
         class MyMetric(Metric):
@@ -2530,45 +2525,38 @@ def test_is_bope_problem(self) -> None:
             )
             self.assertFalse(experiment.is_bope_problem)
 
-    def test_name_and_store_arm_if_not_exists_same_name_different_signature(
-        self,
-    ) -> None:
-        experiment = self.experiment
-        shared_name = "shared_name"
-
-        arm_1 = Arm({"x1": -1.0, "x2": 1.0}, name=shared_name)
-        arm_2 = Arm({"x1": -1.7, "x2": 0.2, "x3": 1})
-        self.assertNotEqual(arm_1.signature, arm_2.signature)
-
-        experiment._register_arm(arm=arm_1)
-        with self.assertRaisesRegex(
-            AxError,
-            f"Arm with name {shared_name} already exists on experiment "
-            f"with different signature.",
-        ):
-            experiment._name_and_store_arm_if_not_exists(
-                arm=arm_2, proposed_name=shared_name
-            )
-
-    def test_name_and_store_arm_if_not_exists_same_proposed_name_different_signature(
-        self,
-    ) -> None:
-        experiment = self.experiment
+    def test_name_and_store_arm_if_not_exists_different_signature(self) -> None:
         shared_name = "shared_name"
-
         arm_1 = Arm({"x1": -1.0, "x2": 1.0}, name=shared_name)
-        arm_2 = Arm({"x1": -1.7, "x2": 0.2, "x3": 1}, name=shared_name)
-        self.assertNotEqual(arm_1.signature, arm_2.signature)
 
-        experiment._register_arm(arm=arm_1)
-        with self.assertRaisesRegex(
-            AxError,
-            f"Arm with name {shared_name} already exists on experiment "
-            f"with different signature.",
-        ):
-            experiment._name_and_store_arm_if_not_exists(
-                arm=arm_2, proposed_name="different proposed name"
-            )
+        cases = [
+            # Unnamed arm with proposed_name matching existing arm's name
+            (
+                "arm_without_name",
+                Arm({"x1": -1.7, "x2": 0.2, "x3": 1}),
+                shared_name,
+            ),
+            # Named arm matching existing arm but with a different proposed_name
+            (
+                "arm_with_shared_name",
+                Arm({"x1": -1.7, "x2": 0.2, "x3": 1}, name=shared_name),
+                "different proposed name",
+            ),
+        ]
+        for label, arm_2, proposed_name in cases:
+            with self.subTest(label=label):
+                self.setUp()
+                experiment = self.experiment
+                experiment._register_arm(arm=arm_1)
+                self.assertNotEqual(arm_1.signature, arm_2.signature)
+                with self.assertRaisesRegex(
+                    AxError,
+                    f"Arm with name {shared_name} already exists on experiment "
+                    f"with different signature.",
+                ):
+                    experiment._name_and_store_arm_if_not_exists(
+                        arm=arm_2, proposed_name=proposed_name
+                    )
 
     def test_sorting_data_by_trial_index_and_arm_name(self) -> None:
         # test sorting data