Consolidate ax/analysis tests (#5007)

Summary:

Part of a 19-diff stack to consolidate repetitive tests across Ax and PTS using `subTest`. Consolidate 11 test files in ax/analysis/ (plotly, healthcheck, graphviz) — adds subTest to top surfaces, metric summary, search space summary, and summary tests.

Differential Revision: D95603642

Author: eonofrey

ax/analysis/graphviz/tests/test_generation_strategy_graph.py

@@ -96,24 +96,23 @@ def setUp(self) -> None:
             ],
         )
 
-    def test_validate_applicable_state_no_gs(self) -> None:
-        """Test that validation fails when no GenerationStrategy is provided."""
+    def test_validate_applicable_state(self) -> None:
+        """Test validation with and without a GenerationStrategy."""
         analysis = GenerationStrategyGraph()
-        result = analysis.validate_applicable_state(
-            experiment=None,
-            generation_strategy=None,
-        )
-        self.assertIsNotNone(result)
-        self.assertIn("requires a GenerationStrategy", result)
-
-    def test_validate_applicable_state_valid(self) -> None:
-        """Test that validation passes with a valid GenerationStrategy."""
-        analysis = GenerationStrategyGraph()
-        result = analysis.validate_applicable_state(
-            experiment=None,
-            generation_strategy=self.node_gs,
-        )
-        self.assertIsNone(result)
+        for label, gs, expect_error in [
+            ("no_gs", None, True),
+            ("valid", self.node_gs, False),
+        ]:
+            with self.subTest(label=label):
+                result = analysis.validate_applicable_state(
+                    experiment=None,
+                    generation_strategy=gs,
+                )
+                if expect_error:
+                    self.assertIsNotNone(result)
+                    self.assertIn("requires a GenerationStrategy", result)
+                else:
+                    self.assertIsNone(result)
 
     def test_compute_step_based_gs(self) -> None:
         """Test computing the graph for a step-based GenerationStrategy."""
@@ -218,26 +217,22 @@ def test_generation_strategy_to_graphviz(self) -> None:
         self.assertIn("Sobol", dot.source)
         self.assertIn("MBM", dot.source)
 
-    def test_add_node_to_graph_current(self) -> None:
-        """Test adding a current node to the graph."""
-        dot = Digraph()
-        node = self.node_gs._nodes[0]
-        _add_node_to_graph(dot=dot, node=node, is_current=True)
-
-        source = dot.source
-        self.assertIn("Sobol", source)
-        self.assertIn("lightblue", source)
-        self.assertIn("bold", source)
-
-    def test_add_node_to_graph_non_current(self) -> None:
-        """Test adding a non-current node to the graph."""
-        dot = Digraph()
-        node = self.node_gs._nodes[1]
-        _add_node_to_graph(dot=dot, node=node, is_current=False)
-
-        source = dot.source
-        self.assertIn("MBM", source)
-        self.assertIn("rounded", source)
+    def test_add_node_to_graph(self) -> None:
+        """Test adding current and non-current nodes to the graph."""
+        for is_current, node_index, expected_name, expected_style in [
+            (True, 0, "Sobol", "lightblue"),
+            (False, 1, "MBM", "rounded"),
+        ]:
+            with self.subTest(is_current=is_current):
+                dot = Digraph()
+                node = self.node_gs._nodes[node_index]
+                _add_node_to_graph(dot=dot, node=node, is_current=is_current)
+
+                source = dot.source
+                self.assertIn(expected_name, source)
+                self.assertIn(expected_style, source)
+                if is_current:
+                    self.assertIn("bold", source)
 
     def test_add_edges_for_node(self) -> None:
         """Test adding edges for a node."""

ax/analysis/graphviz/tests/test_hierarchical_search_space_graph.py

@@ -148,22 +148,17 @@ def test_compute(self) -> None:
                     for dependent_parameter_name in dependent_parameter_names:
                         self.assertIn(dependent_parameter_name, source)
 
-    def test_online(self) -> None:
+    def test_online_and_offline(self) -> None:
         analysis = HierarchicalSearchSpaceGraph()
-        for experiment in get_online_experiments():
-            # If validation fails (i.e. this Experiment is not applicable to HSSGraph)
-            # then skip it in the tests
-            if analysis.validate_applicable_state(experiment=experiment) is not None:
-                continue
-
-            _ = analysis.compute(experiment=experiment)
-
-    def test_offline(self) -> None:
-        analysis = HierarchicalSearchSpaceGraph()
-        for experiment in get_offline_experiments():
-            # If validation fails (i.e. this Experiment is not applicable to HSSGraph)
-            # then skip it in the tests
-            if analysis.validate_applicable_state(experiment=experiment) is not None:
-                continue
-
-            _ = analysis.compute(experiment=experiment)
+        for label, get_experiments in [
+            ("online", get_online_experiments),
+            ("offline", get_offline_experiments),
+        ]:
+            with self.subTest(setting=label):
+                for experiment in get_experiments():
+                    if (
+                        analysis.validate_applicable_state(experiment=experiment)
+                        is not None
+                    ):
+                        continue
+                    _ = analysis.compute(experiment=experiment)

ax/analysis/healthcheck/tests/test_predictable_metrics.py

@@ -212,40 +212,43 @@ def test_random_adapter_not_applicable(self) -> None:
         self.assertIn("RandomAdapter", error)
         self.assertIn("no model to evaluate", error)
 
-    def test_validate_applicable_state_no_experiment(self) -> None:
-        """Test that validation fails when experiment is None."""
+    def test_validate_applicable_state(self) -> None:
+        """Test validation for various input combinations."""
         healthcheck = PredictableMetricsAnalysis()
 
-        error = healthcheck.validate_applicable_state(
-            experiment=None,
-            generation_strategy=self.generation_strategy,
-        )
-
-        self.assertIsNotNone(error)
-        self.assertIn("Experiment", error)
-
-    def test_validate_applicable_state_no_generation_strategy(self) -> None:
-        """Test that validation fails when generation_strategy is None."""
-        healthcheck = PredictableMetricsAnalysis()
-
-        error = healthcheck.validate_applicable_state(
-            experiment=self.experiment,
-            generation_strategy=None,
-        )
-
-        self.assertIsNotNone(error)
-        self.assertIn("GenerationStrategy", error)
-
-    def test_validate_applicable_state_valid_inputs(self) -> None:
-        """Test that validation passes with valid experiment and generation strategy."""
-        healthcheck = PredictableMetricsAnalysis()
-
-        error = healthcheck.validate_applicable_state(
-            experiment=self.experiment,
-            generation_strategy=self.generation_strategy,
-        )
-
-        self.assertIsNone(error)
+        for label, experiment, generation_strategy, expect_error, expected_substr in [
+            (
+                "no_experiment",
+                None,
+                self.generation_strategy,
+                True,
+                "Experiment",
+            ),
+            (
+                "no_generation_strategy",
+                self.experiment,
+                None,
+                True,
+                "GenerationStrategy",
+            ),
+            (
+                "valid_inputs",
+                self.experiment,
+                self.generation_strategy,
+                False,
+                None,
+            ),
+        ]:
+            with self.subTest(label=label):
+                error = healthcheck.validate_applicable_state(
+                    experiment=experiment,
+                    generation_strategy=generation_strategy,
+                )
+                if expect_error:
+                    self.assertIsNotNone(error)
+                    self.assertIn(expected_substr, error)
+                else:
+                    self.assertIsNone(error)
 
     @mock_botorch_optimize
     def test_adapter_resolved_from_generation_strategy(self) -> None:

ax/analysis/healthcheck/tests/test_should_generate_candidates.py

@@ -13,22 +13,17 @@
 
 
 class TestShouldGenerateCandidates(TestCase):
-    def test_should(self) -> None:
-        trial_index = randint(0, 10)
-        card = ShouldGenerateCandidates(
-            should_generate=True,
-            reason="Something reassuring",
-            trial_index=trial_index,
-        ).compute()
-        self.assertEqual(card.get_status(), HealthcheckStatus.PASS)
-        self.assertEqual(card.subtitle, "Something reassuring")
-
-    def test_should_not(self) -> None:
-        trial_index = randint(0, 10)
-        card = ShouldGenerateCandidates(
-            should_generate=False,
-            reason="Something concerning",
-            trial_index=trial_index,
-        ).compute()
-        self.assertEqual(card.get_status(), HealthcheckStatus.WARNING)
-        self.assertEqual(card.subtitle, "Something concerning")
+    def test_should_generate_candidates(self) -> None:
+        for should_generate, reason, expected_status in [
+            (True, "Something reassuring", HealthcheckStatus.PASS),
+            (False, "Something concerning", HealthcheckStatus.WARNING),
+        ]:
+            with self.subTest(should_generate=should_generate):
+                trial_index = randint(0, 10)
+                card = ShouldGenerateCandidates(
+                    should_generate=should_generate,
+                    reason=reason,
+                    trial_index=trial_index,
+                ).compute()
+                self.assertEqual(card.get_status(), expected_status)
+                self.assertEqual(card.subtitle, reason)

ax/analysis/plotly/tests/test_cross_validation.py

@@ -175,64 +175,38 @@ def test_compute_adhoc(self, mock_r2: mock.Mock) -> None:
         )
     )
     @mock_botorch_optimize
-    def test_online(self) -> None:
-        # Test CrossValidationPlot can be computed for a variety of experiments which
-        # resemble those we see in an online setting.
-
-        for experiment in get_online_experiments():
-            for untransform in [True, False]:
-                for refined_metric_name in [None, "foo"]:
-                    generation_strategy = get_default_generation_strategy_at_MBM_node(
-                        experiment=experiment
-                    )
-
-                    # Pick an arbitrary metric from the experiment's optimization config
-                    metric_name = none_throws(
-                        experiment.optimization_config
-                    ).objective.metric_names[0]
-
-                    analysis = CrossValidationPlot(
-                        metric_names=[metric_name],
-                        untransform=untransform,
-                        labels={metric_name: refined_metric_name}
-                        if refined_metric_name
-                        else None,
-                    )
-
-                    _ = analysis.compute(
-                        experiment=experiment, generation_strategy=generation_strategy
-                    )
-
-    @TestCase.ax_long_test(
-        reason=(
-            "cross_validate still too slow under @mock_botorch_optimize for this test"
-        )
-    )
-    @mock_botorch_optimize
-    def test_offline(self) -> None:
-        # Test CrossValidationPlot can be computed for a variety of experiments which
-        # resemble those we see in an online setting.
-
-        for experiment in get_offline_experiments():
-            for untransform in [True, False]:
-                for refined_metric_name in [None, "foo"]:
-                    generation_strategy = get_default_generation_strategy_at_MBM_node(
-                        experiment=experiment
-                    )
-
-                    # Pick an arbitrary metric from the experiment's optimization config
-                    metric_name = none_throws(
-                        experiment.optimization_config
-                    ).objective.metric_names[0]
-
-                    analysis = CrossValidationPlot(
-                        metric_names=[metric_name],
-                        untransform=untransform,
-                        labels={metric_name: refined_metric_name}
-                        if refined_metric_name
-                        else None,
-                    )
-
-                    _ = analysis.compute(
-                        experiment=experiment, generation_strategy=generation_strategy
-                    )
+    def test_online_and_offline(self) -> None:
+        for label, get_experiments in [
+            ("online", get_online_experiments),
+            ("offline", get_offline_experiments),
+        ]:
+            for experiment in get_experiments():
+                for untransform in [True, False]:
+                    for refined_metric_name in [None, "foo"]:
+                        with self.subTest(
+                            setting=label,
+                            untransform=untransform,
+                            refined_metric_name=refined_metric_name,
+                        ):
+                            generation_strategy = (
+                                get_default_generation_strategy_at_MBM_node(
+                                    experiment=experiment
+                                )
+                            )
+
+                            metric_name = none_throws(
+                                experiment.optimization_config
+                            ).objective.metric_names[0]
+
+                            analysis = CrossValidationPlot(
+                                metric_names=[metric_name],
+                                untransform=untransform,
+                                labels={metric_name: refined_metric_name}
+                                if refined_metric_name
+                                else None,
+                            )
+
+                            _ = analysis.compute(
+                                experiment=experiment,
+                                generation_strategy=generation_strategy,
+                            )

ax/analysis/plotly/tests/test_parallel_coordinates.py

@@ -105,30 +105,16 @@ def test_get_parameter_dimension(self) -> None:
             },
         )
 
-    def test_online(self) -> None:
-        # Test ParallelCoordinatesPlot can be computed for a variety of experiments
-        # which resemble those we see in an online setting.
-
-        for experiment in get_online_experiments():
-            analysis = ParallelCoordinatesPlot(
-                # Select and arbitrary metric from the optimization config
-                metric_name=none_throws(
-                    experiment.optimization_config
-                ).objective.metric_names[0]
-            )
-
-            _ = analysis.compute(experiment=experiment)
-
-    def test_offline(self) -> None:
-        # Test ParallelCoordinatesPlot can be computed for a variety of experiments
-        # which resemble those we see in an offline setting.
-
-        for experiment in get_offline_experiments():
-            analysis = ParallelCoordinatesPlot(
-                # Select and arbitrary metric from the optimization config
-                metric_name=none_throws(
-                    experiment.optimization_config
-                ).objective.metric_names[0]
-            )
-
-            _ = analysis.compute(experiment=experiment)
+    def test_online_and_offline(self) -> None:
+        for label, get_experiments in [
+            ("online", get_online_experiments),
+            ("offline", get_offline_experiments),
+        ]:
+            with self.subTest(setting=label):
+                for experiment in get_experiments():
+                    analysis = ParallelCoordinatesPlot(
+                        metric_name=none_throws(
+                            experiment.optimization_config
+                        ).objective.metric_names[0]
+                    )
+                    _ = analysis.compute(experiment=experiment)