ci: establish a hermetic unit-test suite

.github/workflows/pull_request.yml

@@ -7,8 +7,8 @@ on:
       - vdbbench_*
 
 jobs:
-  build:
-    name: Run Python Tests
+  unit-test:
+    name: Hermetic unit tests
     strategy:
       matrix:
         python-version: [3.11, 3.12]
@@ -30,12 +30,12 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -e ".[test]"
+          python -m pip install -e ".[test]"
 
       - name: Run coding checks
         run: |
           make lint
 
-      - name: Test with pytest
+      - name: Run hermetic unit tests
         run: |
-          make unittest
+          make unit-test

Makefile

@@ -1,5 +1,10 @@
-unittest:
-	PYTHONPATH=`pwd` python3 -m pytest tests/test_dataset.py::TestDataSet::test_download_small -svv
+unit-test:
+	python -m pytest --disable-socket tests/unit
+
+unittest: unit-test
+
+e2e-test:
+	python -m pytest tests/e2e -svv
 
 format:
 	PYTHONPATH=`pwd` python3 -m black vectordb_bench

pyproject.toml

@@ -54,6 +54,7 @@ test = [
     "black",
     "ruff",
     "pytest",
+    "pytest-socket",
 ]
 restful         = [ "flask" ]
 qdrant          = [ "qdrant-client" ]

tests/README.md

@@ -0,0 +1,42 @@
+# Test layout
+
+The test suite is being separated by execution requirements.
+
+## Hermetic unit tests
+
+`tests/unit/` is the allowlisted local suite. Tests in this directory must be
+deterministic and must not use network sockets, credentials, live databases,
+containers, or downloaded datasets.
+
+Install only the project and test extra, then run:
+
+```shell
+pip install -e ".[test]"
+make unit-test
+```
+
+`make unittest` is a compatibility alias. The unit target uses
+`pytest-socket` to disable network access and is the suite run by pull-request
+CI.
+
+## Online E2E tests
+
+`tests/e2e/` contains tests that download datasets or require external
+services. Run them explicitly with:
+
+```shell
+make e2e-test
+```
+
+The dataset E2E tests require network access to S3 and/or Aliyun OSS and can
+download large Cohere or LAION datasets. Other E2E tests may require provider
+SDK extras, credentials, containers, or a locally running database; document
+those prerequisites with the test when migrating it.
+
+## Legacy tests
+
+Test modules directly under `tests/` are legacy and are not yet classified.
+They are intentionally outside both Make targets: do not assume they are
+hermetic, and do not move them into `tests/unit/` until they pass with only
+`.[test]` and with network sockets disabled. New tests must go directly into
+`tests/unit/` or `tests/e2e/`.

tests/e2e/test_dataset.py

@@ -0,0 +1,58 @@
+import logging
+import os
+import time
+
+import pytest
+
+from vectordb_bench.backend.data_source import DatasetSource
+from vectordb_bench.backend.dataset import Dataset
+
+log = logging.getLogger("vectordb_bench")
+pytestmark = pytest.mark.integration
+
+
+class TestDataSet:
+    def test_iter_cohere(self):
+        cohere_10m = Dataset.COHERE.manager(10_000_000)
+        cohere_10m.prepare()
+
+        before = time.time()
+        for batch in cohere_10m:
+            log.debug(batch.head(1))
+
+        duration = time.time() - before
+        log.warning("iter through cohere_10m cost=%smin", duration / 60)
+
+    def test_iter_laion(self):
+        laion_100m = Dataset.LAION.manager(100_000_000)
+        laion_100m.prepare(source=DatasetSource.AliyunOSS)
+
+        before = time.time()
+        for batch in laion_100m:
+            log.debug(batch.head(1))
+
+        duration = time.time() - before
+        log.warning("iter through laion_100m cost=%smin", duration / 60)
+
+    def test_download_small(self):
+        openai_50k = Dataset.OPENAI.manager(50_000)
+        files = [
+            "test.parquet",
+            "neighbors.parquet",
+            "neighbors_head_1p.parquet",
+            "neighbors_tail_1p.parquet",
+        ]
+
+        file_path = openai_50k.data_dir.joinpath("test.parquet")
+        DatasetSource.S3.reader().read(
+            openai_50k.data.dir_name.lower(),
+            files=files,
+            local_ds_root=openai_50k.data_dir,
+        )
+
+        os.remove(file_path)
+        DatasetSource.AliyunOSS.reader().read(
+            openai_50k.data.dir_name.lower(),
+            files=files,
+            local_ds_root=openai_50k.data_dir,
+        )

tests/test_dataset.py → tests/unit/test_dataset.py

@@ -8,8 +8,12 @@
 from vectordb_bench import config
 from vectordb_bench.backend import dataset as dataset_module
 from vectordb_bench.backend.clients import MetricType
-from vectordb_bench.backend.data_source import DatasetSource
-from vectordb_bench.backend.dataset import CustomDataset, Dataset, DatasetManager, ParquetDatasetManager
+from vectordb_bench.backend.dataset import (
+    CustomDataset,
+    Dataset,
+    DatasetManager,
+    ParquetDatasetManager,
+)
 from vectordb_bench.backend.filter import LabelFilter, NewIntFilter, non_filter
 
 log = logging.getLogger("vectordb_bench")
@@ -187,8 +191,8 @@ def _write_vector_fixture(data_dir):
 
 class TestDataSet:
     def test_iter_dataset(self):
-        for ds in Dataset:
-            log.info(ds)
+        for dataset in Dataset:
+            log.info(dataset)
 
     def test_cohere(self):
         cohere = Dataset.COHERE.get(100_000)
@@ -201,57 +205,3 @@ def test_cohere(self):
     def test_cohere_error(self):
         with pytest.raises(ValidationError):
             Dataset.COHERE.get(9999)
-
-    def test_iter_cohere(self):
-        cohere_10m = Dataset.COHERE.manager(10_000_000)
-        cohere_10m.prepare()
-
-        import time
-
-        before = time.time()
-        for i in cohere_10m:
-            log.debug(i.head(1))
-
-        dur_iter = time.time() - before
-        log.warning(f"iter through cohere_10m cost={dur_iter/60}min")
-
-    # pytest -sv tests/test_dataset.py::TestDataSet::test_iter_laion
-    def test_iter_laion(self):
-        laion_100m = Dataset.LAION.manager(100_000_000)
-        from vectordb_bench.backend.data_source import DatasetSource
-
-        laion_100m.prepare(source=DatasetSource.AliyunOSS)
-
-        import time
-
-        before = time.time()
-        for i in laion_100m:
-            log.debug(i.head(1))
-
-        dur_iter = time.time() - before
-        log.warning(f"iter through laion_100m cost={dur_iter/60}min")
-
-    def test_download_small(self):
-        openai_50k = Dataset.OPENAI.manager(50_000)
-        files = [
-            "test.parquet",
-            "neighbors.parquet",
-            "neighbors_head_1p.parquet",
-            "neighbors_tail_1p.parquet",
-        ]
-
-        file_path = openai_50k.data_dir.joinpath("test.parquet")
-        import os
-
-        DatasetSource.S3.reader().read(
-            openai_50k.data.dir_name.lower(),
-            files=files,
-            local_ds_root=openai_50k.data_dir,
-        )
-
-        os.remove(file_path)
-        DatasetSource.AliyunOSS.reader().read(
-            openai_50k.data.dir_name.lower(),
-            files=files,
-            local_ds_root=openai_50k.data_dir,
-        )

tests/test_fts_filter_runner.py → tests/unit/test_fts_filter_runner.py

@@ -1,4 +1,5 @@
 import threading
+from types import SimpleNamespace
 
 from vectordb_bench.backend.cases import CaseLabel
 from vectordb_bench.backend.data_source import DatasetSource
@@ -48,7 +49,11 @@ class Case:
         dataset = Dataset()
         filters = filter_obj
 
-    config_obj = type("Config", (), {"stages": [TaskStage.LOAD]})()
+    config_obj = type(
+        "Config",
+        (),
+        {"stages": [TaskStage.LOAD], "case_config": SimpleNamespace(k=10, nq=1)},
+    )()
     runner = CaseRunner.construct(ca=Case(), config=config_obj, dataset_source=DatasetSource.S3)
     init_calls = []
     monkeypatch.setattr(CaseRunner, "init_db", lambda self, drop_old=True: init_calls.append(drop_old))
@@ -70,6 +75,7 @@ class Config:
         db = type("DbConfig", (), {"init_cls": Db})()
         db_config = type("Config", (), {"to_dict": lambda self: {}})()
         db_case_config = object()
+        case_config = SimpleNamespace(k=10)
 
     for filters in (non_filter, NewIntFilter(filter_rate=0.5, int_field="filter_id", int_value=50)):
         case = type(