import importlib import json import sys import types from collections import UserDict from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from core.rag.models.document import Document def _build_fake_vikingdb_modules(): volcengine = types.ModuleType("volcengine") volcengine.__path__ = [] viking_db = types.ModuleType("volcengine.viking_db") class Data(UserDict): def __init__(self, payload): super().__init__(payload) self.fields = payload class DistanceType: L2 = "L2" class IndexType: HNSW = "HNSW" class QuantType: Float = "Float" class FieldType: String = "string" Text = "text" Vector = "vector" class Field: def __init__(self, **kwargs): self.kwargs = kwargs class VectorIndexParams: def __init__(self, **kwargs): self.kwargs = kwargs class _Collection: def __init__(self): self.upsert_data = MagicMock() self.fetch_data = MagicMock(return_value=None) self.delete_data = MagicMock() class _Index: def __init__(self): self.search = MagicMock(return_value=[]) self.search_by_vector = MagicMock(return_value=[]) class VikingDBService: def __init__(self, **kwargs): self.kwargs = kwargs self.create_collection = MagicMock() self.create_index = MagicMock() self.drop_index = MagicMock() self.drop_collection = MagicMock() self._collection = _Collection() self._index = _Index() self.get_collection = MagicMock(return_value=self._collection) self.get_index = MagicMock(return_value=self._index) viking_db.Data = Data viking_db.DistanceType = DistanceType viking_db.Field = Field viking_db.FieldType = FieldType viking_db.IndexType = IndexType viking_db.QuantType = QuantType viking_db.VectorIndexParams = VectorIndexParams viking_db.VikingDBService = VikingDBService return {"volcengine": volcengine, "volcengine.viking_db": viking_db} @pytest.fixture def vikingdb_module(monkeypatch: pytest.MonkeyPatch): for name, module in _build_fake_vikingdb_modules().items(): monkeypatch.setitem(sys.modules, name, module) import dify_vdb_vikingdb.vikingdb_vector as module return importlib.reload(module) def _config(module): return module.VikingDBConfig( access_key="ak", secret_key="sk", host="host", region="region", scheme="https", connection_timeout=10, socket_timeout=20, ) def test_init_get_type_and_has_checks(vikingdb_module): vector = vikingdb_module.VikingDBVector("collection_1", "group-1", _config(vikingdb_module)) assert vector.get_type() == vikingdb_module.VectorType.VIKINGDB assert vector._index_name == "collection_1_idx" assert vector._has_collection() is True assert vector._has_index() is True vector._client.get_collection.side_effect = RuntimeError("missing") assert vector._has_collection() is False vector._client.get_collection.side_effect = None vector._client.get_index.side_effect = RuntimeError("missing") assert vector._has_index() is False def test_create_collection_cache_and_creation_paths(vikingdb_module, monkeypatch: pytest.MonkeyPatch): lock = MagicMock() lock.__enter__.return_value = None lock.__exit__.return_value = None monkeypatch.setattr(vikingdb_module.redis_client, "lock", MagicMock(return_value=lock)) monkeypatch.setattr(vikingdb_module.redis_client, "set", MagicMock()) vector = vikingdb_module.VikingDBVector("collection_1", "group-1", _config(vikingdb_module)) monkeypatch.setattr(vikingdb_module.redis_client, "get", MagicMock(return_value=1)) vector._create_collection(3) vector._client.create_collection.assert_not_called() vector._client.create_index.assert_not_called() monkeypatch.setattr(vikingdb_module.redis_client, "get", MagicMock(return_value=None)) vector._has_collection = MagicMock(return_value=False) vector._has_index = MagicMock(return_value=False) vector._create_collection(4) vector._client.create_collection.assert_called_once() vector._client.create_index.assert_called_once() vikingdb_module.redis_client.set.assert_called_once() def test_create_and_add_texts(vikingdb_module): vector = vikingdb_module.VikingDBVector("collection_1", "group-1", _config(vikingdb_module)) vector._create_collection = MagicMock() vector.add_texts = MagicMock() docs = [Document(page_content="hello", metadata={"doc_id": "id-1"})] vector.create(docs, [[0.1, 0.2]]) vector._create_collection.assert_called_once_with(2) vector.add_texts.assert_called_once_with(docs, [[0.1, 0.2]]) vector = vikingdb_module.VikingDBVector("collection_2", "group-2", _config(vikingdb_module)) docs = [ Document(page_content="a", metadata={"doc_id": "id-a", "document_id": "d-1"}), Document(page_content="b", metadata={"doc_id": "id-b", "document_id": "d-2"}), ] vector.add_texts(docs, [[0.1], [0.2]]) vector._client.get_collection.assert_called() upsert_docs = vector._client.get_collection.return_value.upsert_data.call_args.args[0] assert upsert_docs[0][vikingdb_module.vdb_Field.PRIMARY_KEY] == "id-a" assert upsert_docs[0][vikingdb_module.vdb_Field.GROUP_KEY] == "group-2" def test_text_exists_and_delete_operations(vikingdb_module): vector = vikingdb_module.VikingDBVector("collection_1", "group-1", _config(vikingdb_module)) vector._client.get_collection.return_value.fetch_data.return_value = SimpleNamespace(fields={"message": "ok"}) assert vector.text_exists("id-1") is True vector._client.get_collection.return_value.fetch_data.return_value = SimpleNamespace( fields={"message": "data does not exist"} ) assert vector.text_exists("id-1") is False vector._client.get_collection.return_value.fetch_data.return_value = None assert vector.text_exists("id-1") is False vector.delete_by_ids(["id-1"]) vector._client.get_collection.return_value.delete_data.assert_called_once_with(["id-1"]) vector.get_ids_by_metadata_field = MagicMock(return_value=["id-2"]) vector.delete_by_ids = MagicMock() vector.delete_by_metadata_field("doc_id", "doc-1") vector.delete_by_ids.assert_called_once_with(["id-2"]) def test_get_ids_and_search_helpers(vikingdb_module): vector = vikingdb_module.VikingDBVector("collection_1", "group-1", _config(vikingdb_module)) vector._client.get_index.return_value.search.return_value = [] assert vector.get_ids_by_metadata_field("doc_id", "x") == [] vector._client.get_index.return_value.search.return_value = [ SimpleNamespace(id="a", fields={vikingdb_module.vdb_Field.METADATA_KEY: json.dumps({"doc_id": "x"})}), SimpleNamespace(id="b", fields={vikingdb_module.vdb_Field.METADATA_KEY: json.dumps({"doc_id": "y"})}), SimpleNamespace(id="c", fields={}), ] assert vector.get_ids_by_metadata_field("doc_id", "x") == ["a"] empty_docs = vector._get_search_res([], score_threshold=0.1) assert empty_docs == [] results = [ SimpleNamespace( id="a", score=0.3, fields={ vikingdb_module.vdb_Field.CONTENT_KEY: "doc-a", vikingdb_module.vdb_Field.METADATA_KEY: json.dumps({"document_id": "d-1"}), }, ), SimpleNamespace( id="b", score=0.9, fields={ vikingdb_module.vdb_Field.CONTENT_KEY: "doc-b", vikingdb_module.vdb_Field.METADATA_KEY: json.dumps({"document_id": "d-2"}), }, ), ] docs = vector._get_search_res(results, score_threshold=0.2) assert [doc.page_content for doc in docs] == ["doc-b", "doc-a"] vector._client.get_index.return_value.search_by_vector.return_value = results filtered_docs = vector.search_by_vector([0.1], top_k=2, score_threshold=0.2, document_ids_filter=["d-2"]) assert len(filtered_docs) == 1 assert filtered_docs[0].page_content == "doc-b" assert vector.search_by_full_text("query") == [] def test_delete_drops_index_and_collection_when_present(vikingdb_module): vector = vikingdb_module.VikingDBVector("collection_1", "group-1", _config(vikingdb_module)) vector._has_index = MagicMock(return_value=True) vector._has_collection = MagicMock(return_value=True) vector.delete() vector._client.drop_index.assert_called_once_with("collection_1", "collection_1_idx") vector._client.drop_collection.assert_called_once_with("collection_1") vector._client.drop_index.reset_mock() vector._client.drop_collection.reset_mock() vector._has_index.return_value = False vector._has_collection.return_value = False vector.delete() vector._client.drop_index.assert_not_called() vector._client.drop_collection.assert_not_called() def test_vikingdb_factory_validates_config_and_builds_vector(vikingdb_module, monkeypatch: pytest.MonkeyPatch): factory = vikingdb_module.VikingDBVectorFactory() dataset_with_index = SimpleNamespace( id="dataset-1", index_struct_dict={"vector_store": {"class_prefix": "EXISTING_COLLECTION"}}, index_struct=None, ) dataset_without_index = SimpleNamespace(id="dataset-2", index_struct_dict=None, index_struct=None) monkeypatch.setattr(vikingdb_module.Dataset, "gen_collection_name_by_id", lambda _id: "AUTO_COLLECTION") with patch.object(vikingdb_module, "VikingDBVector", return_value="vector") as vector_cls: monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_ACCESS_KEY", "ak") monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_SECRET_KEY", "sk") monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_HOST", "host") monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_REGION", "region") monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_SCHEME", "https") monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_CONNECTION_TIMEOUT", 10) monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_SOCKET_TIMEOUT", 20) result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock()) result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock()) assert result_1 == "vector" assert result_2 == "vector" assert vector_cls.call_args_list[0].kwargs["collection_name"] == "existing_collection" assert vector_cls.call_args_list[1].kwargs["collection_name"] == "auto_collection" assert dataset_without_index.index_struct is not None @pytest.mark.parametrize( ("field", "message"), [ ("VIKINGDB_ACCESS_KEY", "VIKINGDB_ACCESS_KEY should not be None"), ("VIKINGDB_SECRET_KEY", "VIKINGDB_SECRET_KEY should not be None"), ("VIKINGDB_HOST", "VIKINGDB_HOST should not be None"), ("VIKINGDB_REGION", "VIKINGDB_REGION should not be None"), ("VIKINGDB_SCHEME", "VIKINGDB_SCHEME should not be None"), ], ) def test_vikingdb_factory_raises_when_required_config_missing( vikingdb_module, monkeypatch: pytest.MonkeyPatch, field, message ): factory = vikingdb_module.VikingDBVectorFactory() dataset = SimpleNamespace( id="dataset-1", index_struct_dict={"vector_store": {"class_prefix": "existing"}}, index_struct=None ) monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_ACCESS_KEY", "ak") monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_SECRET_KEY", "sk") monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_HOST", "host") monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_REGION", "region") monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_SCHEME", "https") monkeypatch.setattr(vikingdb_module.dify_config, field, None) with pytest.raises(ValueError, match=message): factory.init_vector(dataset, attributes=[], embeddings=MagicMock())