320 lines
10 KiB
Python
320 lines
10 KiB
Python
"""
|
||
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()
|