Bump dependencies and add a test suite

.github/workflows/ci.yml

@@ -0,0 +1,33 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv and Python
+        uses: astral-sh/setup-uv@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: uv sync --extra dev
+
+      - name: Lint
+        run: uv run ruff check .
+
+      - name: Type check
+        run: uv run ty check
+
+      - name: Test
+        run: uv run pytest -m "not slow"

README.md

@@ -55,6 +55,23 @@ uv pip install --system -e .
 
 Then add `import ucluster.vd.plugin` to your `~/.visidatarc` as described above.
 
+### Testing
+
+Run the test suite with:
+
+```sh
+uv run pytest
+```
+
+The fuzzy-clustering tests run against a deterministic fake encoder, so the
+default run needs no network access. One integration test exercises the real
+sentence-transformer model; it is marked `slow` and skipped by default. Run it
+explicitly (downloads the model on first use) with:
+
+```sh
+uv run pytest -m slow
+```
+
 ## In The Weeds: Architecture & Design
 
 This section will eventually contain a more detailed explanation for how uCluster works. For now, here's a brief overview.

pyproject.toml

@@ -9,16 +9,16 @@ requires-python = ">=3.12"
 dependencies = [
     "loguru>=0.7.3",
     "nltk>=3.9.4",
-    "scikit-learn>=1.8",
-    "sentence-transformers>=5.4",
+    "scikit-learn>=1.9",
+    "sentence-transformers>=5.6",
     "visidata>=3.3",
 ]
 
 [project.optional-dependencies]
 dev = [
-    "pytest>=9.0",
-    "ruff>=0.15",
-    "ty>=0.0.32",
+    "pytest>=9.1",
+    "ruff>=0.15.20",
+    "ty>=0.0.55",
 ]
 
 [build-system]
@@ -28,6 +28,13 @@ build-backend = "hatchling.build"
 [tool.hatch.build.targets.wheel]
 packages = ["ucluster"]
 
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-m 'not slow'"
+markers = [
+    "slow: tests that download or run the real sentence-transformer model",
+]
+
 [tool.ruff]
 line-length = 100
 target-version = "py312"

tests/conftest.py

@@ -0,0 +1,68 @@
+import numpy as np
+import pytest
+from numpy import ndarray
+
+
+class FakeEncoder:
+    """Deterministic stand-in for SentenceTransformer used in fast unit tests.
+
+    Each text is mapped to a vector by keyword: texts containing "alpha" land
+    near one axis, texts containing "bravo" near another, and everything else is
+    pushed to a unique far-away position so it reads as noise. A tiny
+    index-derived jitter keeps points distinct without merging the groups.
+    """
+
+    DIM = 8
+
+    def encode(self, texts: list[str], show_progress_bar: bool = False) -> ndarray:
+        vectors = []
+        for i, text in enumerate(texts):
+            vector = np.zeros(self.DIM, dtype=np.float32)
+            lowered = text.lower()
+            if "alpha" in lowered:
+                vector[0] = 10.0
+                vector[4] = i * 0.001
+            elif "bravo" in lowered:
+                vector[1] = 10.0
+                vector[4] = i * 0.001
+            else:
+                vector[2] = 10.0 + i * 5.0
+            vectors.append(vector)
+        return np.array(vectors, dtype=np.float32)
+
+
+@pytest.fixture
+def fake_encoder() -> FakeEncoder:
+    return FakeEncoder()
+
+
+@pytest.fixture
+def grouped_texts() -> list[str]:
+    """Two clear groups of six plus two outliers, in interleaved order.
+
+    The interleaving ensures tests assert on membership rather than on input
+    position.
+    """
+    alpha = [f"alpha message {n}" for n in range(6)]
+    bravo = [f"bravo message {n}" for n in range(6)]
+    outliers = ["solitary zigzag", "lone quokka"]
+    return [
+        alpha[0], bravo[0], alpha[1], bravo[1], outliers[0],
+        alpha[2], bravo[2], alpha[3], bravo[3], outliers[1],
+        alpha[4], bravo[4], alpha[5], bravo[5],
+    ]
+
+
+@pytest.fixture
+def alpha_indices() -> list[int]:
+    return [0, 2, 5, 7, 10, 12]
+
+
+@pytest.fixture
+def bravo_indices() -> list[int]:
+    return [1, 3, 6, 8, 11, 13]
+
+
+@pytest.fixture
+def outlier_indices() -> list[int]:
+    return [4, 9]

tests/test_exact.py

@@ -0,0 +1,44 @@
+from ucluster import ExactClusterer
+
+
+def _fit(texts: list[str]) -> ExactClusterer:
+    clusterer = ExactClusterer()
+    clusterer.fit(texts)
+    return clusterer
+
+
+def test_exact_duplicates_share_a_cluster():
+    clusters = _fit(["hi there", "hi there", "bye"]).clusters()
+    assert clusters[0] == clusters[1]
+    assert clusters[0] != -1
+    assert clusters[2] == -1
+
+
+def test_matching_is_case_insensitive():
+    clusters = _fit(["Hello", "hello"]).clusters()
+    assert clusters[0] == clusters[1] != -1
+
+
+def test_distinct_duplicate_groups_get_distinct_ids():
+    clusters = _fit(["a", "a", "b", "b"]).clusters()
+    assert clusters[0] == clusters[1]
+    assert clusters[2] == clusters[3]
+    assert clusters[0] != clusters[2]
+
+
+def test_output_length_and_order_match_input():
+    texts = ["a", "b", "a"]
+    clusters = _fit(texts).clusters()
+    assert len(clusters) == len(texts)
+    assert clusters[0] == clusters[2]
+    assert clusters[1] == -1
+
+
+def test_empty_input_yields_empty_clusters():
+    assert _fit([]).clusters() == []
+
+
+def test_probabilities_reflect_membership():
+    clusterer = _fit(["a", "a", "b"])
+    assert clusterer.probabilities() == [1.0, 1.0, 0.0]
+    assert clusterer.outlier_probabilities() == [0.0, 0.0, 1.0]

tests/test_fuzzy.py

@@ -0,0 +1,53 @@
+import pytest
+
+from ucluster import FuzzyClusterer
+
+
+@pytest.fixture
+def fitted(fake_encoder, grouped_texts):
+    clusterer = FuzzyClusterer(encoder=fake_encoder)
+    clusterer.fit(grouped_texts)
+    return clusterer
+
+
+def test_similar_texts_share_a_cluster(fitted, alpha_indices, bravo_indices):
+    clusters = fitted.clusters()
+    alpha_labels = {clusters[i] for i in alpha_indices}
+    bravo_labels = {clusters[i] for i in bravo_indices}
+
+    assert len(alpha_labels) == 1
+    assert len(bravo_labels) == 1
+    assert -1 not in alpha_labels
+    assert -1 not in bravo_labels
+    assert alpha_labels != bravo_labels
+
+
+def test_outliers_are_noise(fitted, outlier_indices):
+    clusters = fitted.clusters()
+    for i in outlier_indices:
+        assert clusters[i] == -1
+
+
+def test_outputs_match_input_length(fitted, grouped_texts):
+    n = len(grouped_texts)
+    assert len(fitted.clusters()) == n
+    assert len(fitted.probabilities()) == n
+    assert len(fitted.outlier_probabilities()) == n
+
+
+def test_probabilities_are_bounded(fitted):
+    assert all(0.0 <= p <= 1.0 for p in fitted.probabilities())
+
+
+def test_outlier_probabilities_complement_probabilities(fitted):
+    probs = fitted.probabilities()
+    outliers = fitted.outlier_probabilities()
+    assert all(abs(o - (1.0 - p)) < 1e-6 for p, o in zip(probs, outliers, strict=True))
+
+
+def test_fit_is_deterministic(fake_encoder, grouped_texts):
+    first = FuzzyClusterer(encoder=fake_encoder)
+    first.fit(grouped_texts)
+    second = FuzzyClusterer(encoder=fake_encoder)
+    second.fit(grouped_texts)
+    assert first.clusters() == second.clusters()

tests/test_fuzzy_real.py

@@ -0,0 +1,65 @@
+"""Opt-in integration test exercising the real sentence-transformer model.
+
+Marked ``slow`` and excluded from default CI runs (``pytest -m "not slow"``).
+Run explicitly with ``pytest -m slow``. The first run downloads ~470MB from the
+Hugging Face Hub; the test skips (rather than fails) when the model cannot be
+loaded, e.g. offline with an empty cache.
+"""
+
+from collections import Counter
+
+import pytest
+from huggingface_hub.errors import HfHubHTTPError, LocalEntryNotFoundError
+
+from ucluster import FuzzyClusterer
+
+pytestmark = pytest.mark.slow
+
+DOG_SENTENCES = [
+    "The dog ran quickly across the park",
+    "A dog sprinted through the park",
+    "The dog dashed across the park grass",
+    "A quick dog raced over the green park",
+]
+
+PASTA_SENTENCES = [
+    "I boiled pasta for dinner tonight",
+    "She cooked spaghetti for the evening meal",
+    "We made pasta for dinner",
+    "He prepared spaghetti this evening",
+]
+
+OUTLIER = "Quantum entanglement continues to puzzle physicists"
+
+
+@pytest.fixture(scope="module")
+def real_clusterer() -> FuzzyClusterer:
+    # Skip only on the specific huggingface_hub failures that mean the model is
+    # genuinely unreachable (offline/uncached or a Hub HTTP error). Any other
+    # exception should surface as a real test failure.
+    try:
+        return FuzzyClusterer()
+    except (LocalEntryNotFoundError, HfHubHTTPError) as exc:
+        pytest.skip(f"sentence-transformer model unavailable: {exc}")
+
+
+def _dominant_label(clusters: list[int], indices: range) -> tuple[int, int]:
+    label, count = Counter(clusters[i] for i in indices).most_common(1)[0]
+    return label, count
+
+
+def test_paraphrase_groups_cluster_apart(real_clusterer: FuzzyClusterer):
+    texts = DOG_SENTENCES + PASTA_SENTENCES + [OUTLIER]
+    real_clusterer.fit(texts)
+    clusters = real_clusterer.clusters()
+
+    dog_label, dog_count = _dominant_label(clusters, range(0, 4))
+    pasta_label, pasta_count = _dominant_label(clusters, range(4, 8))
+
+    # Each paraphrase group should mostly land in a single real cluster...
+    assert dog_label != -1
+    assert pasta_label != -1
+    assert dog_count >= 3
+    assert pasta_count >= 3
+    # ...and the two topics should not share that cluster.
+    assert dog_label != pasta_label

tests/test_preprocess.py

@@ -0,0 +1,18 @@
+from ucluster.text_cluster import preprocess_text
+
+
+def test_lowercases_text():
+    assert preprocess_text("Hello World") == "hello world"
+
+
+def test_separates_punctuation_into_tokens():
+    assert preprocess_text("Hello, world!") == "hello , world !"
+
+
+def test_collapses_runs_of_whitespace():
+    assert preprocess_text("spread   out\t\nwords") == "spread out words"
+
+
+def test_replaces_unencodable_characters_without_raising():
+    result = preprocess_text("caf\udce9")
+    assert isinstance(result, str)

ucluster/text_cluster.py

@@ -1,5 +1,6 @@
 import json
 from collections import defaultdict
+from typing import Protocol
 
 import nltk
 from loguru import logger
@@ -30,20 +31,28 @@ def outlier_probabilities(self) -> list[float]:
         raise NotImplementedError
 
 
+class Encoder(Protocol):
+    def encode(self, texts: list[str], show_progress_bar: bool = ...) -> ndarray: ...
+
+
 class FuzzyClusterer(TextClusterer):
     DEFAULT_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
 
     def __init__(
         self,
         model: str = DEFAULT_MODEL,
+        *,
+        encoder: Encoder | None = None,
         min_cluster_size: int = 3,
         min_samples: int = 3,
         alpha: float = 1.0,
         epsilon: float = 0.0,
     ) -> None:
-        from sentence_transformers import SentenceTransformer
+        if encoder is None:
+            from sentence_transformers import SentenceTransformer
 
-        self._encoder = SentenceTransformer(model)
+            encoder = SentenceTransformer(model)
+        self._encoder = encoder
         self.min_cluster_size = min_cluster_size
         self.min_samples = min_samples
         self.alpha = alpha