fixing tests

Author: mivanit

spd/clustering/merge_history.py

@@ -328,6 +328,19 @@ def normalized(self) -> tuple[MergesArray, dict[str, Any]]:
         # For now, keep using dbg_tensor for overlap_stats analysis
         dbg_tensor(overlap_stats)
 
+        # Convert any Path objects to strings for JSON serialization
+        history_metadatas = []
+        for history in self.data:
+            if history.meta is not None:
+                meta_copy = history.meta.copy()
+                # Convert Path objects to strings
+                for key, value in meta_copy.items():
+                    if isinstance(value, Path):
+                        meta_copy[key] = str(value)
+                history_metadatas.append(meta_copy)
+            else:
+                history_metadatas.append(None)
+
         return (
             merges_array,
             dict(
@@ -336,7 +349,7 @@ def normalized(self) -> tuple[MergesArray, dict[str, Any]]:
                 n_iters=self.n_iters,
                 c_components=c_components,
                 config=self.config.model_dump(mode="json"),
-                history_metadatas=[history.meta for history in self.data],
+                history_metadatas=history_metadatas,
             ),
         )
 

tests/clustering/test_merge_integration.py

@@ -111,7 +111,8 @@ def test_merge_with_popping(self):
 
         # Check results
         assert history is not None
-        assert history.merges.k_groups[0].item() == n_components
+        # First entry is after first merge, so should be n_components - 1
+        assert history.merges.k_groups[0].item() == n_components - 1
         # Final group count depends on pops, but should be less than initial
         assert history.merges.k_groups[-1].item() < n_components
 
@@ -192,8 +193,8 @@ def test_merge_with_small_components(self):
             component_labels=component_labels,
         )
 
-        # Should start with 3 components
-        assert history.merges.k_groups[0].item() == 3
+        # First entry is after first merge, so should be 3 - 1 = 2
+        assert history.merges.k_groups[0].item() == 2
         # Early stopping may occur at 2 groups, so final count could be 2 or 3
         assert history.merges.k_groups[-1].item() >= 2
         assert history.merges.k_groups[-1].item() <= 3

tests/clustering/test_wandb_integration.py

@@ -4,9 +4,10 @@
 from pathlib import Path
 from unittest.mock import Mock, patch
 
-from spd.clustering.merge_history import MergeHistory
-from spd.clustering.s2_clustering import _save_merge_history_to_wandb
-from spd.clustering.s3_normalize_histories import normalize_and_save
+from spd.clustering.merge_config import MergeConfig
+from spd.clustering.merge_history import MergeHistory, MergeHistoryEnsemble
+from spd.clustering.pipeline.s2_clustering import _save_merge_history_to_wandb
+from spd.clustering.pipeline.s3_normalize_histories import normalize_and_save
 
 
 def test_wandb_url_parsing_short_format():
@@ -48,8 +49,6 @@ def test_wandb_url_parsing_short_format():
 
 def test_merge_history_ensemble():
     """Test that MergeHistoryEnsemble can handle multiple histories."""
-    from spd.clustering.merge_config import MergeConfig
-    from spd.clustering.merge_history import MergeHistoryEnsemble
 
     # Create test merge histories
     config = MergeConfig(
@@ -79,7 +78,6 @@ def test_merge_history_ensemble():
 
 def test_save_merge_history_to_wandb():
     """Test that _save_merge_history_to_wandb creates the expected artifact."""
-    from spd.clustering.merge_config import MergeConfig
 
     # Create a real MergeHistory
     config = MergeConfig(
@@ -102,7 +100,7 @@ def test_save_merge_history_to_wandb():
         history_path = Path(tmp_dir) / "test_history.zip"
         history.save(history_path)
 
-        with patch("spd.clustering.s2_clustering.wandb.Artifact") as mock_artifact_class:
+        with patch("spd.clustering.pipeline.s2_clustering.wandb.Artifact") as mock_artifact_class:
             mock_artifact_class.return_value = mock_artifact
 
             # Call the function
@@ -127,7 +125,6 @@ def test_save_merge_history_to_wandb():
 
 def test_wandb_url_field_in_merge_history():
     """Test that MergeHistory can store and serialize wandb_url."""
-    from spd.clustering.merge_config import MergeConfig
 
     # Create a simple config
     config = MergeConfig(
@@ -148,5 +145,5 @@ def test_wandb_url_field_in_merge_history():
         history.save(save_path)
         loaded_history = MergeHistory.read(save_path)
 
-        assert loaded_history
-        assert loaded_history.merges.group_idxs.shape == (0, 5)
+        assert loaded_history is not None
+        assert loaded_history.merges.group_idxs.shape == (10, 5)  # (iters, n_components)