Fixes unit tests in the new ingest pipeline

.github/workflows/pytest.yml

@@ -2,7 +2,9 @@ name: Unit tests
 
 on:
   pull_request:
-    branches: [ main ]
+    branches: 
+      - main
+      - 'feature/new-ingest-process'
   workflow_dispatch:
 
 jobs:

.github/workflows/ruff.yml

@@ -2,7 +2,9 @@ name: Ruff
 
 on:
   pull_request:
-    branches: [ main ]
+    branches: 
+      - main
+      - 'feature/new-ingest-process'
   workflow_dispatch:
 
 jobs:

src/matchbox/common/results.py

@@ -4,7 +4,7 @@
 from collections import defaultdict
 from concurrent.futures import ProcessPoolExecutor
 from enum import StrEnum
-from typing import TYPE_CHECKING, Any, Generic, Hashable, Iterator, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Generic, Hashable, Iterator, TypeVar
 
 import numpy as np
 import pyarrow as pa
@@ -556,14 +556,20 @@ def component_to_hierarchy(table: pa.Table, dtype: pa.DataType = pa.int32) -> pa
 
 
 def to_hierarchical_clusters(
-    probabilities: pa.Table, dtype: pa.DataType = pa.int32
+    probabilities: pa.Table,
+    proc_func: Callable[[pa.Table, pa.DataType], pa.Table] = component_to_hierarchy,
+    dtype: pa.DataType = pa.int32,
+    timeout: int = 300,
 ) -> pa.Table:
     """
     Converts a table of pairwise probabilities into a table of hierarchical clusters.
 
     Args:
         probabilities: Arrow table with columns ['component', 'left', 'right',
             'probability']
+        proc_func: Function to process each component
+        dtype: Arrow data type for parent/child columns
+        timeout: Maximum seconds to wait for each component to process
 
     Returns:
         Arrow table with columns ['parent', 'child', 'probability']
@@ -598,14 +604,19 @@ def to_hierarchical_clusters(
     results = []
     with ProcessPoolExecutor(max_workers=n_cores) as executor:
         futures = [
-            executor.submit(component_to_hierarchy, component_table, dtype)
+            executor.submit(proc_func, component_table, dtype)
             for component_table in component_tables
         ]
 
         for future in futures:
             try:
-                result = future.result()
+                result = future.result(timeout=timeout)
                 results.append(result)
+            except TimeoutError:
+                logic_logger.error(
+                    f"Component processing timed out after {timeout} seconds"
+                )
+                continue
             except Exception as e:
                 logic_logger.error(f"Error processing component: {str(e)}")
                 continue

test/client/test_hierarchy.py

@@ -1,6 +1,8 @@
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
 from functools import lru_cache
 from itertools import chain
-from typing import Any
+from typing import Any, Iterator
 from unittest.mock import patch
 
 import pyarrow as pa
@@ -31,6 +33,23 @@ def _combine_strings(*n: str) -> str:
     return "".join(sorted(letters))
 
 
+@contextmanager
+def parallel_pool_for_tests(
+    max_workers: int = 2, timeout: int = 30
+) -> Iterator[ThreadPoolExecutor]:
+    """Context manager for safe parallel execution in tests using threads.
+
+    Args:
+        max_workers: Maximum number of worker threads
+        timeout: Maximum seconds to wait for each task
+    """
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        try:
+            yield executor
+        finally:
+            executor.shutdown(wait=False, cancel_futures=True)
+
+
 @pytest.mark.parametrize(
     ("parameters"),
     [
@@ -241,11 +260,16 @@ def test_component_to_hierarchy(
                 "probability": [90, 85, 80],
             },
             {
-                ("abc", "a", 90),
-                ("abc", "b", 90),
-                ("abc", "c", 85),
+                ("ab", "a", 90),
+                ("ab", "b", 90),
+                ("bc", "b", 85),
+                ("bc", "c", 85),
+                ("abc", "ab", 85),
+                ("abc", "bc", 85),
+                ("cd", "c", 80),
+                ("cd", "d", 80),
                 ("abcd", "abc", 80),
-                ("abcd", "d", 80),
+                ("abcd", "cd", 80),
             },
         ),
         # Multiple components test case
@@ -257,12 +281,18 @@ def test_component_to_hierarchy(
                 "probability": [90, 85, 95, 92],
             },
             {
-                ("abc", "a", 90),
-                ("abc", "b", 90),
-                ("abc", "c", 85),
-                ("xyz", "x", 95),
-                ("xyz", "y", 95),
-                ("xyz", "z", 92),
+                ("xy", "x", 95),
+                ("xy", "y", 95),
+                ("yz", "y", 92),
+                ("yz", "z", 92),
+                ("xyz", "xy", 92),
+                ("xyz", "yz", 92),
+                ("ab", "a", 90),
+                ("ab", "b", 90),
+                ("bc", "b", 85),
+                ("bc", "c", 85),
+                ("abc", "ab", 85),
+                ("abc", "bc", 85),
             },
         ),
     ],
@@ -304,19 +334,25 @@ def test_hierarchical_clusters(input_data, expected_hierarchy):
         )
     )
 
-    with patch(
-        "matchbox.common.results.combine_integers", side_effect=_combine_strings
+    # Run and compare
+    with (
+        patch(
+            "matchbox.common.results.ProcessPoolExecutor",
+            lambda *args, **kwargs: parallel_pool_for_tests(timeout=30),
+        ),
+        patch("matchbox.common.results.combine_integers", side_effect=_combine_strings),
     ):
-        result = to_hierarchical_clusters(probabilities, dtype=pa.string)
-
-        # Sort result the same way as expected for comparison
-        result = result.sort_by(
-            [
-                ("probability", "descending"),
-                ("parent", "ascending"),
-                ("child", "ascending"),
-            ]
+        result = to_hierarchical_clusters(
+            probabilities, dtype=pa.string, proc_func=component_to_hierarchy
         )
 
-        assert result.schema == expected.schema
-        assert result.equals(expected)
+    result = result.sort_by(
+        [
+            ("probability", "descending"),
+            ("parent", "ascending"),
+            ("child", "ascending"),
+        ]
+    )
+
+    assert result.schema == expected.schema
+    assert result.equals(expected)