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

tests/unit/connectors/test_s3.py

@@ -0,0 +1,598 @@
+"""S3Source 和 S3Destination 测试
+
+测试 S3/MinIO connectors。使用 mock 模式测试，无需真实 S3 服务器或 moto 库。
+
+运行方式：
+    pytest tests/unit/connectors/test_s3.py -v
+"""
+
+import json
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from xparse_client.connectors import S3Destination, S3Source
+from xparse_client.exceptions import DestinationError, SourceError
+
+# ============================================================================
+# Mock 辅助函数
+# ============================================================================
+
+def _setup_s3_mock(mock_get_boto3):
+    """设置 S3 mock 的辅助函数"""
+    mock_boto3 = MagicMock()
+    mock_config_class = MagicMock()
+    mock_client = MagicMock()
+
+    # boto3.client() 返回 mock client
+    mock_boto3.client.return_value = mock_client
+
+    # _get_boto3() 返回 (boto3, Config)
+    mock_get_boto3.return_value = (mock_boto3, mock_config_class)
+
+    return mock_client, mock_boto3, mock_config_class
+
+
+# ============================================================================
+# S3Source 测试
+# ============================================================================
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_initialization(mock_get_boto3):
+    """测试 S3Source 初始化"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+
+    # Mock head_bucket 成功
+    mock_client.head_bucket.return_value = {}
+
+    source = S3Source(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket",
+        prefix="documents/",
+        region="us-west-2"
+    )
+
+    assert source.endpoint == "https://s3.amazonaws.com"
+    assert source.bucket == "test-bucket"
+    assert source.prefix == "documents/"
+    assert source.pattern is None
+    assert source.recursive is False
+
+    # 验证 head_bucket 被调用
+    mock_client.head_bucket.assert_called_once_with(Bucket="test-bucket")
+
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_list_files_basic(mock_get_boto3):
+    """测试列出文件（非递归）"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+
+    # Mock paginator
+    mock_paginator = MagicMock()
+    mock_client.get_paginator.return_value = mock_paginator
+
+    # Mock 分页结果
+    mock_paginator.paginate.return_value = [
+        {
+            "Contents": [
+                {"Key": "file1.pdf"},
+                {"Key": "file2.docx"},
+                {"Key": "subdir/"},  # 目录
+            ]
+        }
+    ]
+
+    source = S3Source(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket",
+        recursive=False
+    )
+
+    files = source.list_files()
+
+    # 应该过滤掉目录
+    assert len(files) == 2
+    assert 'file1.pdf' in files
+    assert 'file2.docx' in files
+
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_list_files_with_prefix(mock_get_boto3):
+    """测试带 prefix 的文件列表"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+
+    mock_paginator = MagicMock()
+    mock_client.get_paginator.return_value = mock_paginator
+
+    mock_paginator.paginate.return_value = [
+        {
+            "Contents": [
+                {"Key": "documents/doc1.pdf"},
+                {"Key": "documents/doc2.docx"},
+            ]
+        }
+    ]
+
+    source = S3Source(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket",
+        prefix="documents/"
+    )
+
+    files = source.list_files()
+
+    # 验证使用了正确的 prefix
+    call_args = mock_paginator.paginate.call_args
+    assert call_args[1]["Prefix"] == "documents/"
+
+    assert len(files) == 2
+
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_list_files_with_pattern(mock_get_boto3):
+    """测试带 pattern 的文件列表"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+
+    mock_paginator = MagicMock()
+    mock_client.get_paginator.return_value = mock_paginator
+
+    mock_paginator.paginate.return_value = [
+        {
+            "Contents": [
+                {"Key": "file1.pdf"},
+                {"Key": "file2.docx"},
+                {"Key": "file3.txt"},
+            ]
+        }
+    ]
+
+    source = S3Source(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket",
+        pattern=["*.pdf"]
+    )
+
+    files = source.list_files()
+
+    assert len(files) == 1
+    assert 'file1.pdf' in files
+
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_list_files_recursive(mock_get_boto3):
+    """测试递归列出文件"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+
+    mock_paginator = MagicMock()
+    mock_client.get_paginator.return_value = mock_paginator
+
+    mock_paginator.paginate.return_value = [
+        {
+            "Contents": [
+                {"Key": "file1.pdf"},
+                {"Key": "subdir/file2.pdf"},
+                {"Key": "subdir/subdir2/file3.pdf"},
+            ]
+        }
+    ]
+
+    source = S3Source(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket",
+        recursive=True
+    )
+
+    files = source.list_files()
+
+    # 递归模式不应该使用 Delimiter
+    call_args = mock_paginator.paginate.call_args
+    assert "Delimiter" not in call_args[1]
+
+    assert len(files) == 3
+
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_read_file(mock_get_boto3):
+    """测试读取文件"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+
+    # Mock get_object
+    test_content = b'S3 file content'
+    mock_body = MagicMock()
+    mock_body.read.return_value = test_content
+
+    mock_client.get_object.return_value = {
+        "Body": mock_body,
+        "ResponseMetadata": {
+            "HTTPHeaders": {
+                "etag": '"abc123"',
+                "last-modified": "Mon, 28 Jan 2024 12:00:00 GMT",
+                "server": "AmazonS3"
+            }
+        }
+    }
+
+    source = S3Source(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket"
+    )
+
+    file_bytes, data_source = source.read_file('test.pdf')
+
+    assert file_bytes == test_content
+    assert data_source['url'] == 's3://test-bucket/test.pdf'
+    assert data_source['version'] == 'abc123'
+    assert data_source['record_locator']['protocol'] == 's3'
+    assert 'date_modified' in data_source
+
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_read_file_not_found(mock_get_boto3):
+    """测试读取不存在的文件"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+
+    # Mock NoSuchKey 异常
+    Exception("NoSuchKey")
+    mock_client.exceptions.NoSuchKey = type('NoSuchKey', (Exception,), {})
+    mock_client.get_object.side_effect = mock_client.exceptions.NoSuchKey("Not found")
+
+    source = S3Source(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket"
+    )
+
+    with pytest.raises(SourceError) as exc_info:
+        source.read_file('nonexistent.pdf')
+
+    assert "S3 对象不存在" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_connection_error(mock_get_boto3):
+    """测试连接错误"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+
+    # Mock head_bucket 失败
+    mock_client.head_bucket.side_effect = Exception("Connection refused")
+
+    with pytest.raises(SourceError) as exc_info:
+        S3Source(
+            endpoint="https://invalid-endpoint.com",
+            access_key="test-key",
+            secret_key="test-secret",
+            bucket="test-bucket"
+        )
+
+    assert "S3 连接失败" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_context_manager(mock_get_boto3):
+    """测试上下文管理器"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+
+    with S3Source(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket"
+    ) as source:
+        assert source is not None
+
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_repr(mock_get_boto3):
+    """测试 __repr__"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+
+    source = S3Source(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket"
+    )
+
+    assert repr(source) == "<S3Source endpoint=https://s3.amazonaws.com bucket=test-bucket>"
+
+
+# ============================================================================
+# S3Destination 测试
+# ============================================================================
+
+@patch('xparse_client.connectors.destinations.s3._get_boto3')
+def test_s3_destination_initialization(mock_get_boto3):
+    """测试 S3Destination 初始化"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+
+    # Mock head_bucket 和 put_object 成功
+    mock_client.head_bucket.return_value = {}
+    mock_client.put_object.return_value = {}
+    mock_client.delete_object.return_value = {}
+
+    dest = S3Destination(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket",
+        prefix="output/"
+    )
+
+    assert dest.endpoint == "https://s3.amazonaws.com"
+    assert dest.bucket == "test-bucket"
+    assert dest.prefix == "output"  # 注意会 strip /
+
+    # 验证测试写入
+    assert mock_client.put_object.called
+
+
+@patch('xparse_client.connectors.destinations.s3._get_boto3')
+def test_s3_destination_write(mock_get_boto3):
+    """测试写入数据"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+    mock_client.put_object.return_value = {}
+    mock_client.delete_object.return_value = {}
+
+    dest = S3Destination(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket"
+    )
+
+    data = [
+        {"element_id": "1", "text": "test1"},
+        {"element_id": "2", "text": "test2"}
+    ]
+    metadata = {"filename": "test.pdf"}
+
+    result = dest.write(data, metadata)
+
+    assert result is True
+
+    # 验证 put_object 被调用（至少 2 次：测试写入 + 实际写入）
+    assert mock_client.put_object.call_count >= 2
+
+    # 检查最后一次调用的参数
+    last_call = mock_client.put_object.call_args_list[-1]
+    assert last_call[1]["Bucket"] == "test-bucket"
+    assert last_call[1]["Key"] == "test.json"
+    assert last_call[1]["ContentType"] == "application/json"
+
+    # 验证写入的数据
+    written_data = json.loads(last_call[1]["Body"].decode("utf-8"))
+    assert len(written_data) == 2
+    assert written_data[0]["element_id"] == "1"
+
+
+@patch('xparse_client.connectors.destinations.s3._get_boto3')
+def test_s3_destination_write_with_prefix(mock_get_boto3):
+    """测试带 prefix 的写入"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+    mock_client.put_object.return_value = {}
+    mock_client.delete_object.return_value = {}
+
+    dest = S3Destination(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket",
+        prefix="output/"
+    )
+
+    data = [{"element_id": "1"}]
+    metadata = {"filename": "test.pdf"}
+
+    result = dest.write(data, metadata)
+
+    assert result is True
+
+    # 检查 key 包含 prefix
+    last_call = mock_client.put_object.call_args_list[-1]
+    assert last_call[1]["Key"] == "output/test.json"
+
+
+@patch('xparse_client.connectors.destinations.s3._get_boto3')
+def test_s3_destination_write_error(mock_get_boto3):
+    """测试写入错误"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+
+    # 第一次 put_object（测试写入）成功
+    mock_client.put_object.side_effect = [{}, Exception("Write failed")]
+    mock_client.delete_object.return_value = {}
+
+    dest = S3Destination(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket"
+    )
+
+    data = [{"element_id": "1"}]
+    metadata = {"filename": "test.pdf"}
+
+    with pytest.raises(DestinationError) as exc_info:
+        dest.write(data, metadata)
+
+    assert "写入 S3 失败" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.destinations.s3._get_boto3')
+def test_s3_destination_connection_error(mock_get_boto3):
+    """测试连接错误"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+
+    # Mock head_bucket 失败
+    mock_client.head_bucket.side_effect = Exception("Bucket not found")
+
+    with pytest.raises(DestinationError) as exc_info:
+        S3Destination(
+            endpoint="https://s3.amazonaws.com",
+            access_key="test-key",
+            secret_key="test-secret",
+            bucket="nonexistent-bucket"
+        )
+
+    assert "S3 连接或写入测试失败" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.destinations.s3._get_boto3')
+def test_s3_destination_context_manager(mock_get_boto3):
+    """测试上下文管理器"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+    mock_client.put_object.return_value = {}
+    mock_client.delete_object.return_value = {}
+
+    with S3Destination(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket"
+    ) as dest:
+        assert dest is not None
+
+
+@patch('xparse_client.connectors.destinations.s3._get_boto3')
+def test_s3_destination_repr(mock_get_boto3):
+    """测试 __repr__"""
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+    mock_client.put_object.return_value = {}
+    mock_client.delete_object.return_value = {}
+
+    dest = S3Destination(
+        endpoint="https://s3.amazonaws.com",
+        access_key="test-key",
+        secret_key="test-secret",
+        bucket="test-bucket"
+    )
+
+    assert repr(dest) == "<S3Destination endpoint=https://s3.amazonaws.com bucket=test-bucket>"
+
+
+# ============================================================================
+# 懒加载失败测试
+# ============================================================================
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_import_error(mock_get_boto3):
+    """测试 boto3 未安装时的 ImportError (Source)"""
+    # Mock ImportError
+    mock_get_boto3.side_effect = ImportError(
+        "使用 S3Source 需要安装 boto3: pip install xparse-client[s3]"
+    )
+
+    with pytest.raises(ImportError) as exc_info:
+        S3Source(
+            endpoint="https://s3.amazonaws.com",
+            bucket="test-bucket",
+            access_key="test-key",
+            secret_key="test-secret"
+        )
+
+    assert "boto3" in str(exc_info.value)
+    assert "xparse-client[s3]" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.destinations.s3._get_boto3')
+def test_s3_destination_import_error(mock_get_boto3):
+    """测试 boto3 未安装时的 ImportError (Destination)"""
+    # Mock ImportError
+    mock_get_boto3.side_effect = ImportError(
+        "使用 S3Destination 需要安装 boto3: pip install xparse-client[s3]"
+    )
+
+    with pytest.raises(ImportError) as exc_info:
+        S3Destination(
+            endpoint="https://s3.amazonaws.com",
+            bucket="test-bucket",
+            access_key="test-key",
+            secret_key="test-secret"
+        )
+
+    assert "boto3" in str(exc_info.value)
+    assert "xparse-client[s3]" in str(exc_info.value)
+
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_list_files_pagination_error(mock_get_boto3):
+    """测试 list_objects_v2 分页失败场景"""
+    from xparse_client.exceptions import SourceError
+
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+
+    # Mock paginator 失败
+    mock_paginator = MagicMock()
+    mock_paginator.paginate.side_effect = Exception("Pagination error")
+    mock_client.get_paginator.return_value = mock_paginator
+
+    source = S3Source(
+        endpoint="https://s3.amazonaws.com",
+        bucket="test-bucket",
+        access_key="test-key",
+        secret_key="test-secret"
+    )
+
+    with pytest.raises(SourceError) as exc_info:
+        list(source.list_files())
+
+    assert "列出文件" in str(exc_info.value) or "list" in str(exc_info.value).lower()
+
+
+@patch('xparse_client.connectors.sources.s3._get_boto3')
+def test_s3_source_read_file_get_object_error(mock_get_boto3):
+    """测试 get_object 失败场景"""
+    from xparse_client.exceptions import SourceError
+
+    mock_client, _, _ = _setup_s3_mock(mock_get_boto3)
+    mock_client.head_bucket.return_value = {}
+
+    # 创建一个 NoSuchKey 异常类
+    class NoSuchKeyError(Exception):
+        pass
+
+    # Mock exceptions 属性
+    mock_exceptions = MagicMock()
+    mock_exceptions.NoSuchKey = NoSuchKeyError
+    mock_client.exceptions = mock_exceptions
+
+    # Mock get_object 抛出 NoSuchKey 异常
+    mock_client.get_object.side_effect = NoSuchKeyError("The specified key does not exist.")
+
+    source = S3Source(
+        endpoint="https://s3.amazonaws.com",
+        bucket="test-bucket",
+        access_key="test-key",
+        secret_key="test-secret"
+    )
+
+    with pytest.raises(SourceError) as exc_info:
+        source.read_file("test.txt")
+
+    assert "不存在" in str(exc_info.value) or "exist" in str(exc_info.value).lower()