wehub-resource-sync
575 lines
19 KiB
Python
575 lines
19 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock, patch, AsyncMock
|
|
import numpy as np
|
|
|
|
pytest.importorskip(
|
|
"qdrant_client",
|
|
reason="qdrant-client is required for Qdrant storage tests",
|
|
)
|
|
|
|
from qdrant_client import models # noqa: E402
|
|
from lightrag.utils import EmbeddingFunc # noqa: E402
|
|
from lightrag.kg.qdrant_impl import QdrantVectorDBStorage # noqa: E402
|
|
|
|
|
|
# Mock QdrantClient
|
|
@pytest.fixture
|
|
def mock_qdrant_client():
|
|
with patch("lightrag.kg.qdrant_impl.QdrantClient") as mock_client_cls:
|
|
client = mock_client_cls.return_value
|
|
client.collection_exists.return_value = False
|
|
client.count.return_value.count = 0
|
|
# Mock payload schema and vector config for get_collection
|
|
collection_info = MagicMock()
|
|
collection_info.payload_schema = {}
|
|
# Mock vector dimension to match mock_embedding_func (768d)
|
|
collection_info.config.params.vectors.size = 768
|
|
client.get_collection.return_value = collection_info
|
|
yield client
|
|
|
|
|
|
# Mock get_data_init_lock to avoid async lock issues in tests
|
|
@pytest.fixture(autouse=True)
|
|
def mock_data_init_lock():
|
|
with patch("lightrag.kg.qdrant_impl.get_data_init_lock") as mock_lock:
|
|
mock_lock_ctx = AsyncMock()
|
|
mock_lock.return_value = mock_lock_ctx
|
|
yield mock_lock
|
|
|
|
|
|
# Mock Embedding function
|
|
@pytest.fixture
|
|
def mock_embedding_func():
|
|
async def embed_func(texts, **kwargs):
|
|
return np.array([[0.1] * 768 for _ in texts])
|
|
|
|
func = EmbeddingFunc(embedding_dim=768, func=embed_func, model_name="test-model")
|
|
return func
|
|
|
|
|
|
async def test_qdrant_collection_naming(mock_qdrant_client, mock_embedding_func):
|
|
"""Test if collection name is correctly generated with model suffix"""
|
|
config = {
|
|
"embedding_batch_num": 10,
|
|
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
|
}
|
|
|
|
storage = QdrantVectorDBStorage(
|
|
namespace="chunks",
|
|
global_config=config,
|
|
embedding_func=mock_embedding_func,
|
|
workspace="test_ws",
|
|
)
|
|
|
|
# Verify collection name contains model suffix
|
|
expected_suffix = "test_model_768d"
|
|
assert expected_suffix in storage.final_namespace
|
|
assert storage.final_namespace == f"lightrag_vdb_chunks_{expected_suffix}"
|
|
|
|
|
|
async def test_qdrant_migration_trigger(mock_qdrant_client, mock_embedding_func):
|
|
"""Test if migration logic is triggered correctly"""
|
|
config = {
|
|
"embedding_batch_num": 10,
|
|
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
|
}
|
|
|
|
storage = QdrantVectorDBStorage(
|
|
namespace="chunks",
|
|
global_config=config,
|
|
embedding_func=mock_embedding_func,
|
|
workspace="test_ws",
|
|
)
|
|
|
|
# Legacy collection name (without model suffix)
|
|
legacy_collection = "lightrag_vdb_chunks"
|
|
|
|
# Setup mocks for migration scenario
|
|
# 1. New collection does not exist, only legacy exists
|
|
mock_qdrant_client.collection_exists.side_effect = lambda name: (
|
|
name == legacy_collection
|
|
)
|
|
|
|
# 2. Legacy collection exists and has data
|
|
migration_state = {"new_workspace_count": 0}
|
|
|
|
def count_mock(collection_name, exact=True, count_filter=None):
|
|
mock_result = MagicMock()
|
|
if collection_name == legacy_collection:
|
|
mock_result.count = 100
|
|
elif collection_name == storage.final_namespace:
|
|
mock_result.count = migration_state["new_workspace_count"]
|
|
else:
|
|
mock_result.count = 0
|
|
return mock_result
|
|
|
|
mock_qdrant_client.count.side_effect = count_mock
|
|
|
|
# 3. Mock scroll for data migration
|
|
mock_point = MagicMock()
|
|
mock_point.id = "old_id"
|
|
mock_point.vector = [0.1] * 768
|
|
mock_point.payload = {"content": "test"} # No workspace_id in payload
|
|
|
|
# When payload_schema is empty, the code first samples payloads to detect workspace_id
|
|
# Then proceeds with migration batches
|
|
# Scroll calls: 1) Sampling (limit=10), 2) Migration batch, 3) End of migration
|
|
mock_qdrant_client.scroll.side_effect = [
|
|
([mock_point], "_"), # Sampling scroll - no workspace_id found
|
|
([mock_point], "next_offset"), # Migration batch
|
|
([], None), # End of migration
|
|
]
|
|
|
|
def upsert_mock(*args, **kwargs):
|
|
migration_state["new_workspace_count"] = 100
|
|
return None
|
|
|
|
mock_qdrant_client.upsert.side_effect = upsert_mock
|
|
|
|
# Initialize storage (triggers migration)
|
|
await storage.initialize()
|
|
|
|
# Verify migration steps
|
|
# 1. Legacy count checked
|
|
mock_qdrant_client.count.assert_any_call(
|
|
collection_name=legacy_collection, exact=True
|
|
)
|
|
|
|
# 2. New collection created
|
|
mock_qdrant_client.create_collection.assert_called()
|
|
|
|
# 3. Data scrolled from legacy
|
|
# First call (index 0) is sampling scroll with limit=10
|
|
# Second call (index 1) is migration batch with limit=500
|
|
assert mock_qdrant_client.scroll.call_count >= 2
|
|
# Check sampling scroll
|
|
sampling_call = mock_qdrant_client.scroll.call_args_list[0]
|
|
assert sampling_call.kwargs["collection_name"] == legacy_collection
|
|
assert sampling_call.kwargs["limit"] == 10
|
|
# Check migration batch scroll
|
|
migration_call = mock_qdrant_client.scroll.call_args_list[1]
|
|
assert migration_call.kwargs["collection_name"] == legacy_collection
|
|
assert migration_call.kwargs["limit"] == 500
|
|
|
|
# 4. Data upserted to new
|
|
mock_qdrant_client.upsert.assert_called()
|
|
|
|
# 5. Payload index created
|
|
mock_qdrant_client.create_payload_index.assert_called()
|
|
|
|
|
|
async def test_qdrant_no_migration_needed(mock_qdrant_client, mock_embedding_func):
|
|
"""Test scenario where new collection already exists (Case 1 in setup_collection)
|
|
|
|
When only the new collection exists and no legacy collection is found,
|
|
the implementation should:
|
|
1. Create payload index on the new collection (ensure index exists)
|
|
2. NOT attempt any data migration (no scroll calls)
|
|
"""
|
|
config = {
|
|
"embedding_batch_num": 10,
|
|
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
|
}
|
|
|
|
storage = QdrantVectorDBStorage(
|
|
namespace="chunks",
|
|
global_config=config,
|
|
embedding_func=mock_embedding_func,
|
|
workspace="test_ws",
|
|
)
|
|
|
|
# Only new collection exists (no legacy collection found)
|
|
mock_qdrant_client.collection_exists.side_effect = lambda name: (
|
|
name == storage.final_namespace
|
|
)
|
|
|
|
# Initialize
|
|
await storage.initialize()
|
|
|
|
# Should create payload index on the new collection (ensure index)
|
|
mock_qdrant_client.create_payload_index.assert_called_with(
|
|
collection_name=storage.final_namespace,
|
|
field_name="workspace_id",
|
|
field_schema=models.KeywordIndexParams(
|
|
type=models.KeywordIndexType.KEYWORD,
|
|
is_tenant=True,
|
|
),
|
|
)
|
|
# Should NOT migrate (no scroll calls since no legacy collection exists)
|
|
mock_qdrant_client.scroll.assert_not_called()
|
|
|
|
|
|
# ============================================================================
|
|
# Tests for scenarios described in design document (Lines 606-649)
|
|
# ============================================================================
|
|
|
|
|
|
async def test_scenario_1_new_workspace_creation(
|
|
mock_qdrant_client, mock_embedding_func
|
|
):
|
|
"""
|
|
场景1：新建workspace
|
|
预期：直接创建lightrag_vdb_chunks_text_embedding_3_large_3072d
|
|
"""
|
|
# Use a large embedding model
|
|
large_model_func = EmbeddingFunc(
|
|
embedding_dim=3072,
|
|
func=mock_embedding_func.func,
|
|
model_name="text-embedding-3-large",
|
|
)
|
|
|
|
config = {
|
|
"embedding_batch_num": 10,
|
|
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
|
}
|
|
|
|
storage = QdrantVectorDBStorage(
|
|
namespace="chunks",
|
|
global_config=config,
|
|
embedding_func=large_model_func,
|
|
workspace="test_new",
|
|
)
|
|
|
|
# Case 3: Neither legacy nor new collection exists
|
|
mock_qdrant_client.collection_exists.return_value = False
|
|
|
|
# Initialize storage
|
|
await storage.initialize()
|
|
|
|
# Verify: Should create new collection with model suffix
|
|
expected_collection = "lightrag_vdb_chunks_text_embedding_3_large_3072d"
|
|
assert storage.final_namespace == expected_collection
|
|
|
|
# Verify create_collection was called with correct name
|
|
create_calls = [
|
|
call for call in mock_qdrant_client.create_collection.call_args_list
|
|
]
|
|
assert len(create_calls) > 0
|
|
assert (
|
|
create_calls[0][0][0] == expected_collection
|
|
or create_calls[0].kwargs.get("collection_name") == expected_collection
|
|
)
|
|
|
|
# Verify no migration was attempted
|
|
mock_qdrant_client.scroll.assert_not_called()
|
|
|
|
print(
|
|
f"✅ Scenario 1: New workspace created with collection '{expected_collection}'"
|
|
)
|
|
|
|
|
|
async def test_scenario_2_legacy_upgrade_migration(
|
|
mock_qdrant_client, mock_embedding_func
|
|
):
|
|
"""
|
|
场景2：从旧版本升级
|
|
已存在lightrag_vdb_chunks（无后缀）
|
|
预期：自动迁移数据到lightrag_vdb_chunks_text_embedding_ada_002_1536d
|
|
注意：迁移后不再自动删除遗留集合，需要手动删除
|
|
"""
|
|
# Use ada-002 model
|
|
ada_func = EmbeddingFunc(
|
|
embedding_dim=1536,
|
|
func=mock_embedding_func.func,
|
|
model_name="text-embedding-ada-002",
|
|
)
|
|
|
|
config = {
|
|
"embedding_batch_num": 10,
|
|
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
|
}
|
|
|
|
storage = QdrantVectorDBStorage(
|
|
namespace="chunks",
|
|
global_config=config,
|
|
embedding_func=ada_func,
|
|
workspace="test_legacy",
|
|
)
|
|
|
|
# Legacy collection name (without model suffix)
|
|
legacy_collection = "lightrag_vdb_chunks"
|
|
new_collection = storage.final_namespace
|
|
|
|
# Case 4: Only legacy collection exists
|
|
mock_qdrant_client.collection_exists.side_effect = lambda name: (
|
|
name == legacy_collection
|
|
)
|
|
|
|
# Mock legacy collection info with 1536d vectors
|
|
legacy_collection_info = MagicMock()
|
|
legacy_collection_info.payload_schema = {}
|
|
legacy_collection_info.config.params.vectors.size = 1536
|
|
mock_qdrant_client.get_collection.return_value = legacy_collection_info
|
|
|
|
migration_state = {"new_workspace_count": 0}
|
|
|
|
def count_mock(collection_name, exact=True, count_filter=None):
|
|
mock_result = MagicMock()
|
|
if collection_name == legacy_collection:
|
|
mock_result.count = 150
|
|
elif collection_name == new_collection:
|
|
mock_result.count = migration_state["new_workspace_count"]
|
|
else:
|
|
mock_result.count = 0
|
|
return mock_result
|
|
|
|
mock_qdrant_client.count.side_effect = count_mock
|
|
|
|
# Mock scroll results (simulate migration in batches)
|
|
mock_points = []
|
|
for i in range(10):
|
|
point = MagicMock()
|
|
point.id = f"legacy-{i}"
|
|
point.vector = [0.1] * 1536
|
|
# No workspace_id in payload - simulates legacy data
|
|
point.payload = {"content": f"Legacy document {i}", "id": f"doc-{i}"}
|
|
mock_points.append(point)
|
|
|
|
# When payload_schema is empty, the code first samples payloads to detect workspace_id
|
|
# Then proceeds with migration batches
|
|
# Scroll calls: 1) Sampling (limit=10), 2) Migration batch, 3) End of migration
|
|
mock_qdrant_client.scroll.side_effect = [
|
|
(mock_points, "_"), # Sampling scroll - no workspace_id found in payloads
|
|
(mock_points, "offset1"), # Migration batch
|
|
([], None), # End of migration
|
|
]
|
|
|
|
def upsert_mock(*args, **kwargs):
|
|
migration_state["new_workspace_count"] = 150
|
|
return None
|
|
|
|
mock_qdrant_client.upsert.side_effect = upsert_mock
|
|
|
|
# Initialize (triggers migration)
|
|
await storage.initialize()
|
|
|
|
# Verify: New collection should be created
|
|
expected_new_collection = "lightrag_vdb_chunks_text_embedding_ada_002_1536d"
|
|
assert storage.final_namespace == expected_new_collection
|
|
|
|
# Verify migration steps
|
|
# 1. Check legacy count
|
|
mock_qdrant_client.count.assert_any_call(
|
|
collection_name=legacy_collection, exact=True
|
|
)
|
|
|
|
# 2. Create new collection
|
|
mock_qdrant_client.create_collection.assert_called()
|
|
|
|
# 3. Scroll legacy data
|
|
scroll_calls = [call for call in mock_qdrant_client.scroll.call_args_list]
|
|
assert len(scroll_calls) >= 1
|
|
assert scroll_calls[0].kwargs["collection_name"] == legacy_collection
|
|
|
|
# 4. Upsert to new collection
|
|
upsert_calls = [call for call in mock_qdrant_client.upsert.call_args_list]
|
|
assert len(upsert_calls) >= 1
|
|
assert upsert_calls[0].kwargs["collection_name"] == new_collection
|
|
|
|
# Note: Legacy collection is NOT automatically deleted after migration
|
|
# Manual deletion is required after data migration verification
|
|
|
|
print(
|
|
f"✅ Scenario 2: Legacy data migrated from '{legacy_collection}' to '{expected_new_collection}'"
|
|
)
|
|
|
|
|
|
async def test_scenario_3_multi_model_coexistence(mock_qdrant_client):
|
|
"""
|
|
场景3：多模型并存
|
|
预期：两个独立的collection，互不干扰
|
|
"""
|
|
|
|
# Model A: bge-small with 768d
|
|
async def embed_func_a(texts, **kwargs):
|
|
return np.array([[0.1] * 768 for _ in texts])
|
|
|
|
model_a_func = EmbeddingFunc(
|
|
embedding_dim=768, func=embed_func_a, model_name="bge-small"
|
|
)
|
|
|
|
# Model B: bge-large with 1024d
|
|
async def embed_func_b(texts, **kwargs):
|
|
return np.array([[0.2] * 1024 for _ in texts])
|
|
|
|
model_b_func = EmbeddingFunc(
|
|
embedding_dim=1024, func=embed_func_b, model_name="bge-large"
|
|
)
|
|
|
|
config = {
|
|
"embedding_batch_num": 10,
|
|
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
|
}
|
|
|
|
# Create storage for workspace A with model A
|
|
storage_a = QdrantVectorDBStorage(
|
|
namespace="chunks",
|
|
global_config=config,
|
|
embedding_func=model_a_func,
|
|
workspace="workspace_a",
|
|
)
|
|
|
|
# Create storage for workspace B with model B
|
|
storage_b = QdrantVectorDBStorage(
|
|
namespace="chunks",
|
|
global_config=config,
|
|
embedding_func=model_b_func,
|
|
workspace="workspace_b",
|
|
)
|
|
|
|
# Verify: Collection names are different
|
|
assert storage_a.final_namespace != storage_b.final_namespace
|
|
|
|
# Verify: Model A collection
|
|
expected_collection_a = "lightrag_vdb_chunks_bge_small_768d"
|
|
assert storage_a.final_namespace == expected_collection_a
|
|
|
|
# Verify: Model B collection
|
|
expected_collection_b = "lightrag_vdb_chunks_bge_large_1024d"
|
|
assert storage_b.final_namespace == expected_collection_b
|
|
|
|
# Verify: Different embedding dimensions are preserved
|
|
assert storage_a.embedding_func.embedding_dim == 768
|
|
assert storage_b.embedding_func.embedding_dim == 1024
|
|
|
|
print("✅ Scenario 3: Multi-model coexistence verified")
|
|
print(f" - Workspace A: {expected_collection_a} (768d)")
|
|
print(f" - Workspace B: {expected_collection_b} (1024d)")
|
|
print(" - Collections are independent")
|
|
|
|
|
|
async def test_case1_empty_legacy_auto_cleanup(mock_qdrant_client, mock_embedding_func):
|
|
"""
|
|
Case 1a: 新旧collection都存在，且旧库为空
|
|
预期：自动删除旧库
|
|
"""
|
|
config = {
|
|
"embedding_batch_num": 10,
|
|
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
|
}
|
|
|
|
storage = QdrantVectorDBStorage(
|
|
namespace="chunks",
|
|
global_config=config,
|
|
embedding_func=mock_embedding_func,
|
|
workspace="test_ws",
|
|
)
|
|
|
|
# Legacy collection name (without model suffix)
|
|
legacy_collection = "lightrag_vdb_chunks"
|
|
new_collection = storage.final_namespace
|
|
|
|
# Mock: Both collections exist
|
|
mock_qdrant_client.collection_exists.side_effect = lambda name: (
|
|
name
|
|
in [
|
|
legacy_collection,
|
|
new_collection,
|
|
]
|
|
)
|
|
|
|
# Mock: Legacy collection is empty (0 records)
|
|
def count_mock(collection_name, exact=True, count_filter=None):
|
|
mock_result = MagicMock()
|
|
if collection_name == legacy_collection:
|
|
mock_result.count = 0 # Empty legacy collection
|
|
else:
|
|
mock_result.count = 100 # New collection has data
|
|
return mock_result
|
|
|
|
mock_qdrant_client.count.side_effect = count_mock
|
|
|
|
# Mock get_collection for Case 2 check
|
|
collection_info = MagicMock()
|
|
collection_info.payload_schema = {"workspace_id": True}
|
|
mock_qdrant_client.get_collection.return_value = collection_info
|
|
|
|
# Initialize storage
|
|
await storage.initialize()
|
|
|
|
# Verify: Empty legacy collection should be automatically cleaned up
|
|
# Empty collections are safe to delete without data loss risk
|
|
delete_calls = [
|
|
call for call in mock_qdrant_client.delete_collection.call_args_list
|
|
]
|
|
assert len(delete_calls) >= 1, "Empty legacy collection should be auto-deleted"
|
|
deleted_collection = (
|
|
delete_calls[0][0][0]
|
|
if delete_calls[0][0]
|
|
else delete_calls[0].kwargs.get("collection_name")
|
|
)
|
|
assert deleted_collection == legacy_collection, (
|
|
f"Expected to delete '{legacy_collection}', but deleted '{deleted_collection}'"
|
|
)
|
|
|
|
print(
|
|
f"✅ Case 1a: Empty legacy collection '{legacy_collection}' auto-deleted successfully"
|
|
)
|
|
|
|
|
|
async def test_case1_nonempty_legacy_warning(mock_qdrant_client, mock_embedding_func):
|
|
"""
|
|
Case 1b: 新旧collection都存在，且旧库有数据
|
|
预期：警告但不删除
|
|
"""
|
|
config = {
|
|
"embedding_batch_num": 10,
|
|
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
|
}
|
|
|
|
storage = QdrantVectorDBStorage(
|
|
namespace="chunks",
|
|
global_config=config,
|
|
embedding_func=mock_embedding_func,
|
|
workspace="test_ws",
|
|
)
|
|
|
|
# Legacy collection name (without model suffix)
|
|
legacy_collection = "lightrag_vdb_chunks"
|
|
new_collection = storage.final_namespace
|
|
|
|
# Mock: Both collections exist
|
|
mock_qdrant_client.collection_exists.side_effect = lambda name: (
|
|
name
|
|
in [
|
|
legacy_collection,
|
|
new_collection,
|
|
]
|
|
)
|
|
|
|
# Mock: Legacy collection has data (50 records)
|
|
def count_mock(collection_name, exact=True, count_filter=None):
|
|
mock_result = MagicMock()
|
|
if collection_name == legacy_collection:
|
|
mock_result.count = 50 # Legacy has data
|
|
else:
|
|
mock_result.count = 100 # New collection has data
|
|
return mock_result
|
|
|
|
mock_qdrant_client.count.side_effect = count_mock
|
|
|
|
# Mock get_collection for Case 2 check
|
|
collection_info = MagicMock()
|
|
collection_info.payload_schema = {"workspace_id": True}
|
|
mock_qdrant_client.get_collection.return_value = collection_info
|
|
|
|
# Initialize storage
|
|
await storage.initialize()
|
|
|
|
# Verify: Legacy collection with data should be preserved
|
|
# We never auto-delete collections that contain data to prevent accidental data loss
|
|
delete_calls = [
|
|
call for call in mock_qdrant_client.delete_collection.call_args_list
|
|
]
|
|
# Check if legacy collection was deleted (it should not be)
|
|
legacy_deleted = any(
|
|
(call[0][0] if call[0] else call.kwargs.get("collection_name"))
|
|
== legacy_collection
|
|
for call in delete_calls
|
|
)
|
|
assert not legacy_deleted, "Legacy collection with data should NOT be auto-deleted"
|
|
|
|
print(
|
|
f"✅ Case 1b: Legacy collection '{legacy_collection}' with data preserved (warning only)"
|
|
)
|