chongming-core/test_worker_storage.py

"""
StorageClient（S3/Minio）单元测试 — 桩替换 minio SDK，不依赖真实 Minio 服务

覆盖:
  - from_config 语义: 未配置 / enabled=false → None；缺必填键 → ValueError；
    auto_create_bucket 确保 bucket；auto_create_bucket=false 且 bucket 不存在 → 降级 None；
    连接失败 → 告警降级为 None
  - 上传/下载/删除/预签名 URL 委托给 minio SDK（桩校验参数与对象数据）
  - Plugin 基类集成: _bootstrap 按 config 的 storage 段初始化 self.storage；
    未配置时保持 None；cleanup 关闭并置 None

运行: uv run python test_worker_storage.py（也兼容 pytest 收集）
"""
import asyncio
import sys
import tempfile
import types
from pathlib import Path

SRC = Path(__file__).parent
sys.path.insert(0, str(SRC))

import chongming_worker.storage as storage_mod  # noqa: E402
from chongming_worker.plugin import Plugin  # noqa: E402
from chongming_worker.storage import StorageClient  # noqa: E402

failures = []


def check(name, cond, detail=""):
    status = "PASS" if cond else "FAIL"
    print(f"[{status}] {name}" + (f" — {detail}" if detail else ""))
    if not cond:
        failures.append(name)


class _WriteResult:
    def __init__(self, object_name):
        self.object_name = object_name


class FakeMinio:
    """最小 minio SDK 桩：记录构造参数与对象操作，bucket/连接行为可控。"""

    instances = []
    # 类级默认行为（每次 from_config 新建实例时生效，测试按需改写）
    default_bucket_exists = True
    default_fail_connect = False

    def __init__(self, endpoint, access_key=None, secret_key=None,
                 secure=False, region=None, http_client=None):
        self.endpoint = endpoint
        self.access_key = access_key
        self.secret_key = secret_key
        self.secure = secure
        self.region = region
        self.http_client = http_client
        self.bucket_exists_ok = FakeMinio.default_bucket_exists
        self.fail_connect = FakeMinio.default_fail_connect
        self.created_buckets = []
        self.objects = {}
        FakeMinio.instances.append(self)

    def bucket_exists(self, bucket):
        if self.fail_connect:
            raise ConnectionError("connection refused")
        return self.bucket_exists_ok

    def make_bucket(self, bucket, location=None):
        self.created_buckets.append((bucket, location))

    def put_object(self, bucket, key, data, length, content_type=None):
        self.objects[key] = data.read(length) if hasattr(data, "read") else data
        return _WriteResult(key)

    def fput_object(self, bucket, key, path, content_type=None):
        self.objects[key] = Path(path).read_bytes()
        return _WriteResult(key)

    def get_object(self, bucket, key):
        return _FakeResponse(self.objects[key])

    def fget_object(self, bucket, key, path):
        Path(path).write_bytes(self.objects[key])

    def remove_object(self, bucket, key):
        self.objects.pop(key, None)

    def presigned_put_object(self, bucket, key, expires=None):
        return f"http://minio/put/{bucket}/{key}?e={expires}"

    def presigned_get_object(self, bucket, key, expires=None):
        return f"http://minio/get/{bucket}/{key}?e={expires}"


class _FakeResponse:
    def __init__(self, data):
        self._data = data
        self.closed = False

    def read(self):
        return self._data

    def close(self):
        self.closed = True

    def release_conn(self):
        pass


def patch_minio():
    """把 sys.modules['minio'] 替换为桩（from minio import Minio 可解析）。"""
    fake_mod = types.ModuleType("minio")
    fake_mod.Minio = FakeMinio
    old = sys.modules.get("minio")
    sys.modules["minio"] = fake_mod
    FakeMinio.instances = []
    FakeMinio.default_bucket_exists = True
    FakeMinio.default_fail_connect = False
    return old


def restore_minio(old):
    if old is None:
        sys.modules.pop("minio", None)
    else:
        sys.modules["minio"] = old


def valid_conf(**overrides):
    conf = {
        "endpoint": "127.0.0.1:9000",
        "access_key": "ak",
        "secret_key": "sk",
        "bucket": "files",
    }
    conf.update(overrides)
    return conf


# ========== 1. from_config 语义 ==========
def test_from_config_disabled():
    check("from_config(None) → None", StorageClient.from_config(None) is None)
    check("from_config({}) → None", StorageClient.from_config({}) is None)
    check(
        "enabled=false → None",
        StorageClient.from_config(valid_conf(enabled=False)) is None,
    )


def test_from_config_missing_keys():
    for key in ("endpoint", "access_key", "secret_key", "bucket"):
        conf = valid_conf()
        conf.pop(key)
        try:
            StorageClient.from_config(conf)
            check(f"缺必填键 {key} → ValueError", False)
        except ValueError as e:
            check(f"缺必填键 {key} → ValueError", key in str(e))


def test_from_config_construct_and_bucket():
    old = patch_minio()
    try:
        # 构造参数正确 + bucket 已存在 → 直接可用
        storage = StorageClient.from_config(valid_conf())
        fake = FakeMinio.instances[-1]
        check(
            "构造参数传递正确",
            storage is not None
            and fake.endpoint == "127.0.0.1:9000"
            and fake.access_key == "ak"
            and fake.secret_key == "sk"
            and fake.secure is False
            and storage.bucket == "files"
            and storage.endpoint == "127.0.0.1:9000",
        )

        # secure=true / region 透传
        storage = StorageClient.from_config(
            valid_conf(secure=True, region="cn-east")
        )
        fake = FakeMinio.instances[-1]
        check(
            "secure/region 透传",
            storage is not None
            and fake.secure is True
            and fake.region == "cn-east",
        )

        # bucket 不存在 + auto_create_bucket（默认 true）→ make_bucket
        FakeMinio.default_bucket_exists = False
        storage = StorageClient.from_config(valid_conf())
        fake = FakeMinio.instances[-1]
        check(
            "bucket 不存在时自动创建",
            storage is not None and fake.created_buckets == [("files", None)],
            f"created_buckets={fake.created_buckets}",
        )

        # bucket 不存在 + auto_create_bucket=false → 降级 None
        storage = StorageClient.from_config(valid_conf(auto_create_bucket=False))
        check("auto_create_bucket=false 且 bucket 缺失 → None", storage is None)

        # 连接失败 → 降级 None（不抛错）
        FakeMinio.default_fail_connect = True
        storage = StorageClient.from_config(valid_conf())
        check("连接失败 → 降级 None", storage is None)
    finally:
        restore_minio(old)


# ========== 2. 对象操作委托（同步 + async） ==========
def test_object_operations():
    old = patch_minio()
    try:
        storage = StorageClient.from_config(valid_conf())
        fake = FakeMinio.instances[-1]

        # 同步: put_bytes / get_bytes 往返
        key1 = storage.put_bytes("a/1.txt", b"hello")
        check("put_bytes 上传并返回 object_key", key1 == "a/1.txt" and fake.objects == {"a/1.txt": b"hello"})
        check("get_bytes 下载往返", storage.get_bytes("a/1.txt") == b"hello")

        # 同步: put_file / get_file 往返
        with tempfile.TemporaryDirectory() as tmp:
            src = Path(tmp) / "src.bin"
            src.write_bytes(b"file-data")
            storage.put_file("a/2.bin", src)
            check("put_file 上传", fake.objects.get("a/2.bin") == b"file-data")
            dst = storage.get_file("a/2.bin", Path(tmp) / "dst.bin")
            check("get_file 下载", dst.read_bytes() == b"file-data")

        # 预签名 URL
        put_url = storage.presigned_put_url("a/3.bin")
        get_url = storage.presigned_get_url("a/3.bin")
        check(
            "预签名 PUT/GET URL",
            "http://minio/put/files/a/3.bin" in put_url
            and "http://minio/get/files/a/3.bin" in get_url,
            f"put={put_url}, get={get_url}",
        )

        # 删除
        storage.delete("a/1.txt")
        check("delete 删除对象", "a/1.txt" not in fake.objects)

        # async 包装（to_thread → 同步核心）
        async def run():
            await storage.aput_bytes("a/async.bin", b"async-data")
            data = await storage.aget_bytes("a/async.bin")
            url = await storage.apresigned_put_url("a/async.bin")
            await storage.adelete("a/async.bin")
            return data, url, "a/async.bin" not in fake.objects

        data, url, deleted = asyncio.run(run())
        check(
            "async 包装可用（aput/aget/apresign/adelete）",
            data == b"async-data" and "http://minio/put" in url and deleted,
        )
    finally:
        restore_minio(old)


# ========== 3. Plugin 基类集成 ==========
class _StoragePlugin(Plugin):
    type = "storage_test"
    db_name = None
    auto_bind = False


def test_plugin_integration():
    old = patch_minio()
    try:
        async def run_with_storage():
            p = _StoragePlugin()
            p.config = {"storage": valid_conf()}
            await p._bootstrap()
            ok_init = p.storage is not None
            ok_bucket = p.storage.bucket == "files"
            await p.cleanup("worker/storage_test/x")
            ok_cleanup = p.storage is None
            return ok_init, ok_bucket, ok_cleanup

        ok_init, ok_bucket, ok_cleanup = asyncio.run(run_with_storage())
        check("Plugin._bootstrap 按 storage 配置初始化 self.storage", ok_init and ok_bucket)
        check("Plugin.cleanup 关闭并置 None", ok_cleanup)
    finally:
        restore_minio(old)

    async def run_without_storage():
        p = _StoragePlugin()
        p.config = {}  # 无 storage 段
        await p._bootstrap()
        return p.storage

    val = asyncio.run(run_without_storage())
    check("无 storage 配置 → self.storage 保持 None", val is None)


def main():
    print("========== StorageClient（S3/Minio）单元测试 ==========")
    for fn in (
        test_from_config_disabled,
        test_from_config_missing_keys,
        test_from_config_construct_and_bucket,
        test_object_operations,
        test_plugin_integration,
    ):
        fn()
    print("\n" + ("=" * 40))
    if failures:
        print(f"❌ {len(failures)} 项未通过: {failures}")
        sys.exit(1)
    print("✅ StorageClient 全部通过")


if __name__ == "__main__":
    main()