chore: import upstream snapshot with attribution

tests/test_embedding_batch_size.py

@@ -0,0 +1,69 @@
+"""Tests for embedding batch size and CPU thread configuration."""
+
+from unittest.mock import patch
+
+from leann.embedding_compute import (
+    _cap_cuda_batch_by_vram,
+    _parse_positive_int_env,
+    _resolve_adaptive_batch_size,
+    _resolve_cpu_thread_count,
+)
+
+
+def test_parse_positive_int_env_default(monkeypatch):
+    monkeypatch.delenv("LEANN_TEST_INT", raising=False)
+    assert _parse_positive_int_env("LEANN_TEST_INT", 256) == 256
+
+
+def test_parse_positive_int_env_override(monkeypatch):
+    monkeypatch.setenv("LEANN_TEST_INT", "32")
+    assert _parse_positive_int_env("LEANN_TEST_INT", 256) == 32
+
+
+def test_parse_positive_int_env_invalid(monkeypatch):
+    monkeypatch.setenv("LEANN_TEST_INT", "not-a-number")
+    assert _parse_positive_int_env("LEANN_TEST_INT", 256) == 256
+
+
+def test_resolve_adaptive_batch_size_cuda(monkeypatch):
+    monkeypatch.setenv("LEANN_CUDA_BATCH_SIZE", "64")
+    assert _resolve_adaptive_batch_size("cuda", "BAAI/bge-base-en-v1.5") == 64
+
+
+def test_resolve_adaptive_batch_size_mps_qwen(monkeypatch):
+    monkeypatch.delenv("LEANN_MPS_BATCH_SIZE", raising=False)
+    assert _resolve_adaptive_batch_size("mps", "Qwen/Qwen3-Embedding-0.6B") == 32
+
+
+def test_resolve_cpu_threads(monkeypatch):
+    monkeypatch.setenv("LEANN_CPU_THREADS", "16")
+    assert _resolve_cpu_thread_count() == 16
+
+
+def test_cap_cuda_batch_by_vram_disabled(monkeypatch):
+    monkeypatch.setenv("LEANN_CUDA_AUTO_BATCH", "0")
+    with patch("torch.cuda.is_available", return_value=True):
+        with patch("torch.cuda.mem_get_info", return_value=(100, 1000)):
+            assert _cap_cuda_batch_by_vram(256) == 256
+
+
+def test_cap_cuda_batch_by_vram_small_gpu(monkeypatch):
+    monkeypatch.delenv("LEANN_CUDA_AUTO_BATCH", raising=False)
+    # Typical free VRAM on a 4 GiB GPU after loading a base-sized encoder.
+    one_gb = 1024**3
+    with patch("torch.cuda.is_available", return_value=True):
+        with patch("torch.cuda.mem_get_info", return_value=(one_gb, 4 * one_gb)):
+            capped = _cap_cuda_batch_by_vram(256, max_length=512)
+    assert capped < 256
+    assert capped >= 1
+
+
+def test_cap_cuda_batch_by_vram_four_gb_gpu(monkeypatch):
+    """Regression: 4 GiB RTX A1000 reports ~3.2 GiB free; cap should land near 76."""
+    monkeypatch.delenv("LEANN_CUDA_AUTO_BATCH", raising=False)
+    free_vram = int(3.2 * 1024**3)
+    with patch("torch.cuda.is_available", return_value=True):
+        with patch("torch.cuda.mem_get_info", return_value=(free_vram, 4 * 1024**3)):
+            capped = _cap_cuda_batch_by_vram(256, max_length=512)
+    assert capped <= 85
+    assert capped >= 1