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

tests/unit/connectors/test_local_connectors.py

@@ -0,0 +1,324 @@
+"""LocalSource 和 LocalDestination 集成测试
+
+测试本地文件系统 Source 和 Destination，包括：
+- LocalSource: 文件列表、文件读取、pattern 匹配、递归搜索
+- LocalDestination: 文件写入、元数据处理、中间结果
+- 错误处理：目录不存在、文件不存在、权限错误
+"""
+
+import json
+
+import pytest
+
+from xparse_client.connectors import LocalDestination, LocalSource
+from xparse_client.exceptions import SourceError
+
+# ============================================================================
+# LocalSource 测试
+# ============================================================================
+
+
+def test_local_source_initialization(tmp_path):
+    """测试 LocalSource 正确初始化"""
+    source = LocalSource(directory=str(tmp_path))
+
+    assert source.directory == tmp_path
+    assert source.pattern is None  # None 表示匹配所有文件
+    assert source.recursive is False
+
+
+def test_local_source_directory_not_exist():
+    """测试目录不存在时抛出异常"""
+    with pytest.raises(SourceError) as exc_info:
+        LocalSource(directory="/nonexistent/directory")
+
+    assert "目录不存在" in str(exc_info.value)
+    assert exc_info.value.connector_type == "local"
+
+
+def test_local_source_not_a_directory(tmp_path):
+    """测试路径不是目录时抛出异常"""
+    # 创建一个文件
+    file_path = tmp_path / "file.txt"
+    file_path.write_text("test")
+
+    with pytest.raises(SourceError) as exc_info:
+        LocalSource(directory=str(file_path))
+
+    assert "不是目录" in str(exc_info.value)
+
+
+def test_local_source_list_files_empty(tmp_path):
+    """测试列出空目录"""
+    source = LocalSource(directory=str(tmp_path))
+    files = source.list_files()
+
+    assert files == []
+
+
+def test_local_source_list_files_basic(tmp_path):
+    """测试列出文件（非递归）"""
+    # 创建测试文件
+    (tmp_path / "file1.pdf").write_text("test1")
+    (tmp_path / "file2.docx").write_text("test2")
+    (tmp_path / "file3.txt").write_text("test3")
+
+    # 创建子目录文件（不应该被列出）
+    subdir = tmp_path / "subdir"
+    subdir.mkdir()
+    (subdir / "file4.pdf").write_text("test4")
+
+    source = LocalSource(directory=str(tmp_path))
+    files = source.list_files()
+
+    # 不递归，只列出顶层文件
+    assert len(files) == 3
+    assert "file1.pdf" in files
+    assert "file2.docx" in files
+    assert "file3.txt" in files
+    assert "subdir/file4.pdf" not in files
+
+
+def test_local_source_list_files_with_pattern(tmp_path):
+    """测试 pattern 过滤"""
+    # 创建测试文件
+    (tmp_path / "file1.pdf").write_text("test1")
+    (tmp_path / "file2.docx").write_text("test2")
+    (tmp_path / "file3.txt").write_text("test3")
+
+    source = LocalSource(directory=str(tmp_path), pattern=["*.pdf", "*.docx"])
+    files = source.list_files()
+
+    assert len(files) == 2
+    assert "file1.pdf" in files
+    assert "file2.docx" in files
+    assert "file3.txt" not in files
+
+
+def test_local_source_list_files_recursive(tmp_path):
+    """测试递归列出文件"""
+    # 创建测试文件
+    (tmp_path / "file1.pdf").write_text("test1")
+
+    # 创建子目录和文件
+    subdir = tmp_path / "subdir"
+    subdir.mkdir()
+    (subdir / "file2.pdf").write_text("test2")
+
+    subdir2 = subdir / "subdir2"
+    subdir2.mkdir()
+    (subdir2 / "file3.pdf").write_text("test3")
+
+    source = LocalSource(directory=str(tmp_path), recursive=True)
+    files = source.list_files()
+
+    assert len(files) == 3
+    assert "file1.pdf" in files
+    assert "subdir/file2.pdf" in files
+    assert "subdir/subdir2/file3.pdf" in files
+
+
+def test_local_source_read_file(tmp_path):
+    """测试读取文件"""
+    # 创建测试文件
+    test_file = tmp_path / "test.pdf"
+    test_content = b"PDF content"
+    test_file.write_bytes(test_content)
+
+    source = LocalSource(directory=str(tmp_path))
+    file_bytes, data_source = source.read_file("test.pdf")
+
+    assert file_bytes == test_content
+    assert "url" in data_source
+    assert "version" in data_source
+    assert "date_created" in data_source
+    assert "date_modified" in data_source
+    assert data_source["record_locator"]["protocol"] == "file"
+
+
+def test_local_source_read_file_not_found(tmp_path):
+    """测试读取不存在的文件"""
+    source = LocalSource(directory=str(tmp_path))
+
+    with pytest.raises(SourceError) as exc_info:
+        source.read_file("nonexistent.pdf")
+
+    assert "不存在" in str(exc_info.value)
+    assert exc_info.value.connector_type == "local"
+
+
+def test_local_source_read_file_path_traversal(tmp_path):
+    """测试路径遍历攻击防护"""
+    source = LocalSource(directory=str(tmp_path))
+
+    with pytest.raises(SourceError) as exc_info:
+        source.read_file("../../../etc/passwd")
+
+    assert "非法路径" in str(exc_info.value)
+
+
+def test_local_source_context_manager(tmp_path):
+    """测试上下文管理器"""
+    with LocalSource(directory=str(tmp_path)) as source:
+        assert source is not None
+
+
+# ============================================================================
+# LocalDestination 测试
+# ============================================================================
+
+
+def test_local_destination_initialization(tmp_path):
+    """测试 LocalDestination 正确初始化"""
+    output_dir = tmp_path / "output"
+    dest = LocalDestination(output_dir=str(output_dir))
+
+    assert dest.output_dir == output_dir
+    assert output_dir.exists()
+
+
+def test_local_destination_create_nested_directory(tmp_path):
+    """测试创建嵌套目录"""
+    output_dir = tmp_path / "level1" / "level2" / "output"
+    dest = LocalDestination(output_dir=str(output_dir))
+
+    assert dest.output_dir.exists()
+
+
+def test_local_destination_write_basic(tmp_path):
+    """测试写入基础数据"""
+    dest = LocalDestination(output_dir=str(tmp_path))
+
+    data = [
+        {"element_id": "1", "text": "test1"},
+        {"element_id": "2", "text": "test2"},
+    ]
+    metadata = {"filename": "test.pdf"}
+
+    result = dest.write(data, metadata)
+
+    assert result is True
+
+    # 验证文件创建
+    output_file = tmp_path / "test.json"
+    assert output_file.exists()
+
+    # 验证文件内容
+    with open(output_file, encoding="utf-8") as f:
+        written_data = json.load(f)
+    assert written_data == data
+
+
+def test_local_destination_write_with_stage(tmp_path):
+    """测试写入中间结果（带 stage）"""
+    dest = LocalDestination(output_dir=str(tmp_path))
+
+    data = [{"element_id": "1", "text": "parsed text"}]
+    metadata = {"filename": "test.pdf", "stage": "parse"}
+
+    result = dest.write(data, metadata)
+
+    assert result is True
+
+    # 验证文件名包含 stage
+    output_file = tmp_path / "test_parse.json"
+    assert output_file.exists()
+
+
+def test_local_destination_write_without_filename(tmp_path):
+    """测试写入时没有 filename（使用默认值）"""
+    dest = LocalDestination(output_dir=str(tmp_path))
+
+    data = [{"element_id": "1"}]
+    metadata = {}
+
+    result = dest.write(data, metadata)
+
+    assert result is True
+
+    # 验证使用默认文件名
+    output_file = tmp_path / "output.json"
+    assert output_file.exists()
+
+
+def test_local_destination_context_manager(tmp_path):
+    """测试上下文管理器"""
+    with LocalDestination(output_dir=str(tmp_path)) as dest:
+        assert dest is not None
+
+
+# ============================================================================
+# 集成测试
+# ============================================================================
+
+
+def test_local_source_and_destination_integration(tmp_path):
+    """测试 LocalSource 和 LocalDestination 集成"""
+    # 设置输入目录
+    input_dir = tmp_path / "input"
+    input_dir.mkdir()
+    (input_dir / "file1.pdf").write_text("test1")
+    (input_dir / "file2.pdf").write_text("test2")
+
+    # 设置输出目录
+    output_dir = tmp_path / "output"
+
+    # 创建 source 和 destination
+    source = LocalSource(directory=str(input_dir), pattern=["*.pdf"])
+    dest = LocalDestination(output_dir=str(output_dir))
+
+    # 读取并写入
+    files = source.list_files()
+    assert len(files) == 2
+
+    for file_path in files:
+        content, data_source = source.read_file(file_path)
+
+        # 模拟处理结果
+        elements = [
+            {
+                "element_id": "1",
+                "text": content.decode("utf-8"),
+                "data_source": data_source
+            }
+        ]
+
+        dest.write(elements, {"filename": file_path})
+
+    # 验证输出
+    output_files = list(output_dir.glob("*.json"))
+    assert len(output_files) == 2
+
+
+def test_local_source_read_file_permission_error(tmp_path):
+    """测试 read_file 遇到权限错误"""
+    from xparse_client.exceptions import SourceError
+
+    # 创建测试文件
+    test_file = tmp_path / "test.txt"
+    test_file.write_text("test content")
+
+    source = LocalSource(directory=str(tmp_path))
+
+    # Mock open 抛出 PermissionError
+    from unittest.mock import patch
+    with patch('builtins.open', side_effect=PermissionError("Permission denied")):
+        with pytest.raises(SourceError) as exc_info:
+            source.read_file("test.txt")
+
+        assert "无权限读取文件" in str(exc_info.value) or "Permission" in str(exc_info.value)
+
+
+def test_local_source_read_file_general_error(tmp_path):
+    """测试 read_file 遇到一般错误"""
+    from xparse_client.exceptions import SourceError
+
+    source = LocalSource(directory=str(tmp_path))
+
+    # Mock open 抛出一般 Exception
+    from unittest.mock import patch
+    with patch('builtins.open', side_effect=Exception("Unknown error")):
+        with pytest.raises(SourceError) as exc_info:
+            source.read_file("test.txt")
+
+        assert "读取文件失败" in str(exc_info.value) or "Unknown" in str(exc_info.value)

tests/unit/connectors/test_milvus.py

@@ -0,0 +1,368 @@
+"""MilvusDestination 测试
+
+测试 Milvus 向量数据库 Destination。使用 mock 模式测试，无需真实 Milvus 实例。
+
+运行方式：
+    pytest tests/unit/connectors/test_milvus.py -v
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from xparse_client.connectors import MilvusDestination
+from xparse_client.exceptions import DestinationError
+
+# ============================================================================
+# Mock 辅助函数
+# ============================================================================
+
+def _setup_milvus_mock(mock_get_pymilvus):
+    """设置 Milvus mock 的辅助函数"""
+    mock_milvus_client_class = MagicMock()
+    mock_data_type = MagicMock()
+    mock_client_instance = MagicMock()
+
+    # MilvusClient() 返回 mock 实例
+    mock_milvus_client_class.return_value = mock_client_instance
+
+    # _get_pymilvus() 返回 (MilvusClient, DataType)
+    mock_get_pymilvus.return_value = (mock_milvus_client_class, mock_data_type)
+
+    return mock_client_instance, mock_milvus_client_class, mock_data_type
+
+
+# ============================================================================
+# MilvusDestination 测试
+# ============================================================================
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_destination_initialization(mock_get_pymilvus):
+    """测试 MilvusDestination 初始化"""
+    mock_client, _, _ = _setup_milvus_mock(mock_get_pymilvus)
+
+    # Mock list_collections
+    mock_client.list_collections.return_value = []
+
+    dest = MilvusDestination(
+        db_path="milvus.db",
+        collection_name="test_collection",
+        dimension=1024
+    )
+
+    assert dest.db_path == "milvus.db"
+    assert dest.collection_name == "test_collection"
+    assert dest.dimension == 1024
+
+    # collection 不存在时应该创建
+    # 但我们 mock 的 list_collections 返回空，所以需要检查创建逻辑
+    # 实际实现中，不存在时会创建
+
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_destination_with_api_key(mock_get_pymilvus):
+    """测试使用 API key 初始化"""
+    mock_client, mock_class, _ = _setup_milvus_mock(mock_get_pymilvus)
+    mock_client.list_collections.return_value = []
+
+    MilvusDestination(
+        db_path="https://zilliz-cloud.com",
+        collection_name="test_collection",
+        dimension=1024,
+        api_key="test-api-key"
+    )
+
+    # 验证 API key 被传递
+    call_args = mock_class.call_args
+    assert call_args[1]["token"] == "test-api-key"
+
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_destination_collection_exists(mock_get_pymilvus):
+    """测试 collection 已存在的情况"""
+    mock_client, _, _ = _setup_milvus_mock(mock_get_pymilvus)
+
+    # Mock collection 已存在
+    mock_client.list_collections.return_value = ["test_collection"]
+
+    MilvusDestination(
+        db_path="milvus.db",
+        collection_name="test_collection",
+        dimension=1024
+    )
+
+    # 应该不创建新的 collection
+    assert not mock_client.create_collection.called
+
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_destination_write_with_embeddings(mock_get_pymilvus):
+    """测试写入带 embeddings 的数据"""
+    mock_client, _, _ = _setup_milvus_mock(mock_get_pymilvus)
+    mock_client.list_collections.return_value = ["test_collection"]
+
+    dest = MilvusDestination(
+        db_path="milvus.db",
+        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
+
+    # 验证 insert 被调用
+    mock_client.insert.assert_called_once()
+    call_args = mock_client.insert.call_args
+
+    # 验证 collection_name
+    assert call_args[1]["collection_name"] == "test_collection"
+
+    # 验证插入的数据包含必需字段
+    inserted_data = call_args[1]["data"]
+    assert len(inserted_data) == 1
+    assert inserted_data[0]["element_id"] == "elem_001"
+    assert inserted_data[0]["record_id"] == "rec_123"
+    assert len(inserted_data[0]["embeddings"]) == 1024
+
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_destination_write_missing_embeddings(mock_get_pymilvus):
+    """测试写入缺少 embeddings 的数据（应该返回 False）"""
+    mock_client, _, _ = _setup_milvus_mock(mock_get_pymilvus)
+    mock_client.list_collections.return_value = ["test_collection"]
+
+    dest = MilvusDestination(
+        db_path="milvus.db",
+        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
+
+    # insert 不应该被调用
+    assert not mock_client.insert.called
+
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_destination_write_wrong_dimension(mock_get_pymilvus):
+    """测试写入错误维度的向量（仍然会尝试写入）"""
+    mock_client, _, _ = _setup_milvus_mock(mock_get_pymilvus)
+    mock_client.list_collections.return_value = ["test_collection"]
+
+    dest = MilvusDestination(
+        db_path="milvus.db",
+        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
+
+    # insert 会被调用
+    assert mock_client.insert.called
+
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_destination_batch_write(mock_get_pymilvus):
+    """测试批量写入"""
+    mock_client, _, _ = _setup_milvus_mock(mock_get_pymilvus)
+    mock_client.list_collections.return_value = ["test_collection"]
+
+    dest = MilvusDestination(
+        db_path="milvus.db",
+        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
+
+    # 验证 insert 被调用
+    mock_client.insert.assert_called_once()
+    call_args = mock_client.insert.call_args
+    inserted_data = call_args[1]["data"]
+    assert len(inserted_data) == 10
+
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_destination_connection_error(mock_get_pymilvus):
+    """测试连接错误"""
+    mock_client, mock_class, _ = _setup_milvus_mock(mock_get_pymilvus)
+
+    # Mock 客户端创建失败
+    mock_class.side_effect = Exception("Connection failed")
+
+    with pytest.raises(DestinationError) as exc_info:
+        MilvusDestination(
+            db_path="invalid-path",
+            collection_name="test_collection",
+            dimension=1024
+        )
+
+    assert "初始化" in str(exc_info.value) or "连接" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_destination_write_error(mock_get_pymilvus):
+    """测试写入错误"""
+    mock_client, _, _ = _setup_milvus_mock(mock_get_pymilvus)
+    mock_client.list_collections.return_value = ["test_collection"]
+
+    # Mock insert 失败
+    mock_client.insert.side_effect = Exception("Write failed")
+
+    dest = MilvusDestination(
+        db_path="milvus.db",
+        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.milvus._get_pymilvus')
+def test_milvus_destination_context_manager(mock_get_pymilvus):
+    """测试上下文管理器"""
+    mock_client, _, _ = _setup_milvus_mock(mock_get_pymilvus)
+    mock_client.list_collections.return_value = []
+
+    with MilvusDestination(
+        db_path="milvus.db",
+        collection_name="test_collection",
+        dimension=1024
+    ) as dest:
+        assert dest is not None
+
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_destination_repr(mock_get_pymilvus):
+    """测试 __repr__"""
+    mock_client, _, _ = _setup_milvus_mock(mock_get_pymilvus)
+    mock_client.list_collections.return_value = []
+
+    dest = MilvusDestination(
+        db_path="milvus.db",
+        collection_name="test_collection",
+        dimension=1024
+    )
+
+    assert repr(dest) == "<MilvusDestination db_path=milvus.db collection=test_collection>"
+
+
+# ============================================================================
+# 懒加载失败测试
+# ============================================================================
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_import_error(mock_get_pymilvus):
+    """测试 pymilvus 未安装时的 ImportError"""
+    # Mock ImportError
+    mock_get_pymilvus.side_effect = ImportError(
+        "使用 MilvusDestination 需要安装 pymilvus: pip install xparse-client[milvus]"
+    )
+
+    with pytest.raises(ImportError) as exc_info:
+        MilvusDestination(
+            db_path="milvus.db",
+            collection_name="test_collection",
+            dimension=1024
+        )
+
+    assert "pymilvus" in str(exc_info.value)
+    assert "xparse-client[milvus]" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_collection_creation_error(mock_get_pymilvus):
+    """测试集合创建失败的场景"""
+    mock_client, _, mock_data_type = _setup_milvus_mock(mock_get_pymilvus)
+
+    # Mock collection 不存在
+    mock_client.has_collection.return_value = False
+
+    # Mock create_collection 失败
+    mock_client.create_collection.side_effect = Exception("Schema error")
+
+    with pytest.raises(DestinationError) as exc_info:
+        MilvusDestination(
+            db_path="milvus.db",
+            collection_name="test_collection",
+            dimension=1024
+        )
+
+    assert "初始化" in str(exc_info.value) or "创建" in str(exc_info.value) or "Schema" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.destinations.milvus._get_pymilvus')
+def test_milvus_list_collections_error(mock_get_pymilvus):
+    """测试 has_collection 失败的场景"""
+    mock_client, _, _ = _setup_milvus_mock(mock_get_pymilvus)
+
+    # Mock has_collection 失败
+    mock_client.has_collection.side_effect = Exception("Connection error")
+
+    with pytest.raises(DestinationError) as exc_info:
+        MilvusDestination(
+            db_path="milvus.db",
+            collection_name="test_collection",
+            dimension=1024
+        )
+
+    assert "初始化" in str(exc_info.value) or "连接" in str(exc_info.value) or "Connection" in str(exc_info.value)