xparse-client 0.2.11__py3-none-any.whl → 0.3.0b3__py3-none-any.whl

tests/unit/connectors/test_qdrant.py

@@ -0,0 +1,399 @@
+"""QdrantDestination 测试
+
+测试 Qdrant 向量数据库 Destination。使用 mock 模式测试，无需真实 Qdrant 实例。
+
+运行方式：
+    pytest tests/unit/connectors/test_qdrant.py -v
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from xparse_client.connectors import QdrantDestination
+from xparse_client.exceptions import DestinationError
+
+# ============================================================================
+# Mock 辅助函数
+# ============================================================================
+
+def _setup_qdrant_mock(mock_get_qdrant):
+    """设置 Qdrant mock 的辅助函数"""
+    mock_qdrant_client_class = MagicMock()
+    mock_distance = MagicMock()
+    mock_payload_schema = MagicMock()
+    mock_point_struct = MagicMock()
+    mock_vector_params = MagicMock()
+    mock_client_instance = MagicMock()
+
+    # QdrantClient() 返回 mock 实例
+    mock_qdrant_client_class.return_value = mock_client_instance
+
+    # _get_qdrant() 返回所有需要的类
+    mock_get_qdrant.return_value = (
+        mock_qdrant_client_class,
+        mock_distance,
+        mock_payload_schema,
+        mock_point_struct,
+        mock_vector_params,
+    )
+
+    return (
+        mock_client_instance,
+        mock_qdrant_client_class,
+        mock_distance,
+        mock_payload_schema,
+        mock_point_struct,
+        mock_vector_params,
+    )
+
+
+# ============================================================================
+# QdrantDestination 测试
+# ============================================================================
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_destination_initialization(mock_get_qdrant):
+    """测试 QdrantDestination 初始化"""
+    mock_client, _, _, _, _, _ = _setup_qdrant_mock(mock_get_qdrant)
+
+    # Mock collection 不存在
+    mock_client.collection_exists.return_value = False
+
+    dest = QdrantDestination(
+        url="http://localhost:6333",
+        collection_name="test_collection",
+        dimension=1024
+    )
+
+    assert dest.url == "http://localhost:6333"
+    assert dest.collection_name == "test_collection"
+    assert dest.dimension == 1024
+
+    # 验证 collection 被创建
+    assert mock_client.create_collection.called
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_destination_with_api_key(mock_get_qdrant):
+    """测试使用 API key 初始化"""
+    mock_client, mock_class, _, _, _, _ = _setup_qdrant_mock(mock_get_qdrant)
+    mock_client.collection_exists.return_value = False
+
+    QdrantDestination(
+        url="http://localhost:6333",
+        collection_name="test_collection",
+        dimension=1024,
+        api_key="test-api-key"
+    )
+
+    # 验证 API key 被传递
+    call_args = mock_class.call_args
+    assert call_args[1]["api_key"] == "test-api-key"
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_destination_prefer_grpc(mock_get_qdrant):
+    """测试使用 gRPC"""
+    mock_client, mock_class, _, _, _, _ = _setup_qdrant_mock(mock_get_qdrant)
+    mock_client.collection_exists.return_value = False
+
+    QdrantDestination(
+        url="http://localhost:6333",
+        collection_name="test_collection",
+        dimension=1024,
+        prefer_grpc=True
+    )
+
+    # 验证 prefer_grpc 被传递
+    call_args = mock_class.call_args
+    assert call_args[1]["prefer_grpc"] is True
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_destination_collection_exists(mock_get_qdrant):
+    """测试 collection 已存在的情况"""
+    mock_client, _, _, _, _, _ = _setup_qdrant_mock(mock_get_qdrant)
+
+    # Mock collection 已存在
+    mock_client.collection_exists.return_value = True
+
+    QdrantDestination(
+        url="http://localhost:6333",
+        collection_name="test_collection",
+        dimension=1024
+    )
+
+    # collection 存在时不应该创建，但需要设置 payload schema
+    # 实际实现中 collection_exists=True 后会跳过创建但设置 schema
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_destination_write_with_embeddings(mock_get_qdrant):
+    """测试写入带 embeddings 的数据"""
+    mock_client, _, _, _, mock_point_struct, _ = _setup_qdrant_mock(mock_get_qdrant)
+    mock_client.collection_exists.return_value = True
+
+    dest = QdrantDestination(
+        url="http://localhost:6333",
+        collection_name="test_collection",
+        dimension=1024
+    )
+
+    data = [
+        {
+            "element_id": "elem_001",
+            "text": "Test text",
+            "embeddings": [0.1] * 1024,
+            "metadata": {"page": 1}
+        }
+    ]
+    metadata = {"filename": "test.pdf", "record_id": "rec_123"}
+
+    result = dest.write(data, metadata)
+
+    assert result is True
+
+    # 验证 upsert 被调用
+    mock_client.upsert.assert_called_once()
+    call_args = mock_client.upsert.call_args
+
+    # 验证 collection_name
+    assert call_args[1]["collection_name"] == "test_collection"
+
+    # 验证使用了 PointStruct
+    assert mock_point_struct.called
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_destination_write_missing_embeddings(mock_get_qdrant):
+    """测试写入缺少 embeddings 的数据（应该返回 False）"""
+    mock_client, _, _, _, _, _ = _setup_qdrant_mock(mock_get_qdrant)
+    mock_client.collection_exists.return_value = True
+
+    dest = QdrantDestination(
+        url="http://localhost:6333",
+        collection_name="test_collection",
+        dimension=1024
+    )
+
+    # 缺少 embeddings 的数据
+    data = [
+        {
+            "element_id": "elem_001",
+            "text": "Test text"
+        }
+    ]
+    metadata = {"filename": "test.pdf", "record_id": "rec_123"}
+
+    # 应该返回 False（因为没有有效数据）
+    result = dest.write(data, metadata)
+    assert result is False
+
+    # upsert 不应该被调用
+    assert not mock_client.upsert.called
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_destination_write_wrong_dimension(mock_get_qdrant):
+    """测试写入错误维度的向量（仍然会尝试写入）"""
+    mock_client, _, _, _, _, _ = _setup_qdrant_mock(mock_get_qdrant)
+    mock_client.collection_exists.return_value = True
+
+    dest = QdrantDestination(
+        url="http://localhost:6333",
+        collection_name="test_collection",
+        dimension=1024
+    )
+
+    # 错误维度的 embeddings
+    data = [
+        {
+            "element_id": "elem_001",
+            "text": "Test text",
+            "embeddings": [0.1] * 512  # 应该是 1024
+        }
+    ]
+    metadata = {"filename": "test.pdf", "record_id": "rec_123"}
+
+    # 实现中会尝试写入（由数据库层面检查维度）
+    result = dest.write(data, metadata)
+    assert result is True
+
+    # upsert 会被调用
+    assert mock_client.upsert.called
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_destination_batch_write(mock_get_qdrant):
+    """测试批量写入"""
+    mock_client, _, _, _, mock_point_struct, _ = _setup_qdrant_mock(mock_get_qdrant)
+    mock_client.collection_exists.return_value = True
+
+    dest = QdrantDestination(
+        url="http://localhost:6333",
+        collection_name="test_collection",
+        dimension=1024
+    )
+
+    # 多条数据
+    data = [
+        {
+            "element_id": f"elem_{i:03d}",
+            "text": f"Text {i}",
+            "embeddings": [0.1 * i] * 1024
+        }
+        for i in range(10)
+    ]
+    metadata = {"filename": "test.pdf", "record_id": "rec_123"}
+
+    result = dest.write(data, metadata)
+
+    assert result is True
+
+    # 验证 upsert 被调用
+    mock_client.upsert.assert_called_once()
+
+    # 验证创建了 10 个 Point
+    assert mock_point_struct.call_count == 10
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_destination_connection_error(mock_get_qdrant):
+    """测试连接错误"""
+    mock_client, mock_class, _, _, _, _ = _setup_qdrant_mock(mock_get_qdrant)
+
+    # Mock 客户端创建失败
+    mock_class.side_effect = Exception("Connection failed")
+
+    with pytest.raises(DestinationError) as exc_info:
+        QdrantDestination(
+            url="http://invalid-host:6333",
+            collection_name="test_collection",
+            dimension=1024
+        )
+
+    assert "初始化" in str(exc_info.value) or "连接" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_destination_write_error(mock_get_qdrant):
+    """测试写入错误"""
+    mock_client, _, _, _, _, _ = _setup_qdrant_mock(mock_get_qdrant)
+    mock_client.collection_exists.return_value = True
+
+    # Mock upsert 失败
+    mock_client.upsert.side_effect = Exception("Write failed")
+
+    dest = QdrantDestination(
+        url="http://localhost:6333",
+        collection_name="test_collection",
+        dimension=1024
+    )
+
+    data = [
+        {
+            "element_id": "elem_001",
+            "text": "Test",
+            "embeddings": [0.1] * 1024
+        }
+    ]
+    metadata = {"filename": "test.pdf", "record_id": "rec_123"}
+
+    with pytest.raises(DestinationError) as exc_info:
+        dest.write(data, metadata)
+
+    assert "写入" in str(exc_info.value) or "Write" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_destination_context_manager(mock_get_qdrant):
+    """测试上下文管理器"""
+    mock_client, _, _, _, _, _ = _setup_qdrant_mock(mock_get_qdrant)
+    mock_client.collection_exists.return_value = False
+
+    with QdrantDestination(
+        url="http://localhost:6333",
+        collection_name="test_collection",
+        dimension=1024
+    ) as dest:
+        assert dest is not None
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_destination_repr(mock_get_qdrant):
+    """测试 __repr__"""
+    mock_client, _, _, _, _, _ = _setup_qdrant_mock(mock_get_qdrant)
+    mock_client.collection_exists.return_value = False
+
+    dest = QdrantDestination(
+        url="http://localhost:6333",
+        collection_name="test_collection",
+        dimension=1024
+    )
+
+    assert repr(dest) == "<QdrantDestination url=http://localhost:6333 collection=test_collection>"
+
+
+# ============================================================================
+# 懒加载失败测试
+# ============================================================================
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_import_error(mock_get_qdrant):
+    """测试 qdrant-client 未安装时的 ImportError"""
+    # Mock ImportError
+    mock_get_qdrant.side_effect = ImportError(
+        "使用 QdrantDestination 需要安装 qdrant-client: pip install xparse-client[qdrant]"
+    )
+
+    with pytest.raises(ImportError) as exc_info:
+        QdrantDestination(
+            url="http://localhost:6333",
+            collection_name="test_collection",
+            dimension=1024
+        )
+
+    assert "qdrant" in str(exc_info.value)
+    assert "xparse-client[qdrant]" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_collection_creation_error(mock_get_qdrant):
+    """测试集合创建失败的场景"""
+    mock_client, _, _, _, mock_vector_params, _ = _setup_qdrant_mock(mock_get_qdrant)
+
+    # Mock collection 不存在
+    mock_client.collection_exists.return_value = False
+
+    # Mock create_collection 失败
+    mock_client.create_collection.side_effect = Exception("Schema error")
+
+    with pytest.raises(DestinationError) as exc_info:
+        QdrantDestination(
+            url="http://localhost:6333",
+            collection_name="test_collection",
+            dimension=1024
+        )
+
+    # 错误信息应该包含"连接"或"Schema"
+    assert "连接" in str(exc_info.value) or "Schema" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.destinations.qdrant._get_qdrant')
+def test_qdrant_collection_exists_error(mock_get_qdrant):
+    """测试 get_collections 检查失败的场景"""
+    mock_client, _, _, _, _, _ = _setup_qdrant_mock(mock_get_qdrant)
+
+    # Mock get_collections 失败
+    mock_client.get_collections.side_effect = Exception("Connection error")
+
+    with pytest.raises(DestinationError) as exc_info:
+        QdrantDestination(
+            url="http://localhost:6333",
+            collection_name="test_collection",
+            dimension=1024
+        )
+
+    # 错误信息应该包含"连接"或"Connection"
+    assert "连接" in str(exc_info.value) or "Connection" in str(exc_info.value)