Fix 3D comparison plots and exercise headless rendering

pyppur/utils/visualization.py

@@ -273,7 +273,12 @@ def plot_comparison(
     if embedding_dim not in (2, 3):
         raise ValueError(f"Can only plot 2D or 3D embeddings, got {embedding_dim}D")
 
-    fig, axes = plt.subplots(1, n_plots, figsize=figsize)
+    fig, axes = plt.subplots(
+        1,
+        n_plots,
+        figsize=figsize,
+        subplot_kw={"projection": "3d"} if embedding_dim == 3 else None,
+    )
 
     if n_plots == 1:
         axes = [axes]

tests/test_utils_extended.py

@@ -86,19 +86,30 @@ def test_plot_reconstruction():
     plt.close("all")
 
 
-@pytest.mark.skip("Visualization test - may be flaky in CI")
-def test_plot_comparison_2d():
-    """Test comparison plot with 2D embeddings."""
-
-
-@pytest.mark.skip("Visualization test - may be flaky in CI")
-def test_plot_comparison_3d():
-    """Test comparison plot with 3D embeddings."""
-
-
-@pytest.mark.skip("Visualization test - may be flaky in CI")
-def test_plot_comparison_without_labels():
-    """Test comparison plot without labels."""
+@pytest.mark.parametrize("dimension", [2, 3])
+@pytest.mark.parametrize("with_labels", [False, True])
+def test_plot_comparison_renders_each_embedding(dimension, with_labels, tmp_path):
+    """Comparison plots render every embedding on the requested axes."""
+    rng = np.random.default_rng(42)
+    embeddings = {
+        name: rng.normal(size=(12, dimension)) for name in ("first", "second")
+    }
+    labels = np.arange(12) % 3 if with_labels else None
+
+    fig = plot_comparison(embeddings, labels, title="Comparison")
+    try:
+        axes = [ax for ax in fig.axes if ax.get_title() in embeddings]
+        assert len(axes) == 2
+        assert all(
+            ax.name == ("3d" if dimension == 3 else "rectilinear") for ax in axes
+        )
+        for ax in axes:
+            assert len(ax.collections) == 1
+        output = tmp_path / "comparison.png"
+        fig.savefig(output)
+        assert output.stat().st_size > 0
+    finally:
+        plt.close(fig)
 
 
 def test_visualization_error_cases():