242 lines
9.9 KiB
Python
242 lines
9.9 KiB
Python
"""
|
||
[workers.config] 两种配置模式（toml / cache）语义单测 —— 无需网关与缓存后端。
|
||
|
||
用 dict 作缓存桩，验证：
|
||
- cache 模式：以缓存为准（插件初始化前解析即读到缓存值），缓存缺失才用 toml 引导；
|
||
- toml 模式：以 config.toml 为准，启动时覆盖缓存，使前端（读缓存）与运行值一致；
|
||
- count：toml 模式用文件值覆盖缓存，cache 模式缓存有值用缓存、缺失用文件值引导；
|
||
- get_config 的来源标注（值来自缓存 vs 回退读 toml 种子）如实区分。
|
||
|
||
对应改动：chongming_worker.runtime.resolve_runtime_config / resolve_startup_config /
|
||
_read_worker_config_section_with_source（Plugin.create_workers 在 on_init 前调用）。
|
||
"""
|
||
import os
|
||
import shutil
|
||
import sys
|
||
import tempfile
|
||
|
||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../utils/worker/src"))
|
||
|
||
import chongming_worker.runtime as r
|
||
|
||
PASS = 0
|
||
FAIL = 0
|
||
|
||
WT = "demo"
|
||
TOML_CFG = {"api_key": "TOML-KEY", "model": "toml-model", "heartbeat": {"interval": 5}}
|
||
|
||
|
||
def report(name, cond, detail=""):
|
||
global PASS, FAIL
|
||
if cond:
|
||
PASS += 1
|
||
print(f"✅ {name}" + (f" — {detail}" if detail else ""))
|
||
else:
|
||
FAIL += 1
|
||
print(f"❌ {name}" + (f" — {detail}" if detail else ""))
|
||
|
||
|
||
def write_plugin_toml(root: str) -> None:
|
||
"""在工作目录下写 config/workers/<type>/config.toml（_plugin_config_path 相对 cwd）。"""
|
||
d = os.path.join(root, "config", "workers", WT)
|
||
os.makedirs(d, exist_ok=True)
|
||
with open(os.path.join(d, "config.toml"), "w", encoding="utf-8") as f:
|
||
f.write(
|
||
"[workers]\n"
|
||
f'type = "{WT}"\n'
|
||
'module = "plugins.demo"\n'
|
||
"count = 3\n"
|
||
"\n"
|
||
"[workers.config]\n"
|
||
'api_key = "TOML-KEY"\n'
|
||
'model = "toml-model"\n'
|
||
"\n"
|
||
"[workers.config.heartbeat]\n"
|
||
"interval = 5\n"
|
||
)
|
||
|
||
|
||
def run() -> int:
|
||
workdir = tempfile.mkdtemp(prefix="cm_cfg_mode_")
|
||
cwd = os.getcwd()
|
||
try:
|
||
write_plugin_toml(workdir)
|
||
os.chdir(workdir)
|
||
|
||
key = r._configs_cache_key(WT)
|
||
# _discover_plugins 展平出来的形态：注册字段 + [workers.config] 各键平铺到顶层
|
||
base = {
|
||
"type": WT, "count": 3,
|
||
"api_key": "TOML-KEY", "model": "toml-model",
|
||
"heartbeat": {"interval": 5},
|
||
}
|
||
cached_cfg = {"api_key": "CACHE-KEY", "model": "cache-model"}
|
||
|
||
# ── 1. cache 模式：以缓存为准（插件必须读到缓存值）──────────────
|
||
cache = {key: dict(cached_cfg)}
|
||
got = r.resolve_runtime_config(WT, base, cache, "cache")
|
||
report(
|
||
"cache 模式：返回值取缓存值（覆盖 toml 同键）",
|
||
got.get("api_key") == "CACHE-KEY" and got.get("model") == "cache-model",
|
||
f"api_key={got.get('api_key')}",
|
||
)
|
||
report(
|
||
"cache 模式：toml 独有键保留（缓存缺该键时兜底）",
|
||
got.get("heartbeat") == {"interval": 5},
|
||
f"heartbeat={got.get('heartbeat')}",
|
||
)
|
||
report(
|
||
"cache 模式：不覆盖既有缓存值",
|
||
cache[key] == cached_cfg,
|
||
f"cache={cache[key]}",
|
||
)
|
||
|
||
# 插件初始化路径（process.py 先 set_worker_cache/set_worker_mode 再 create_workers）
|
||
r.set_worker_cache(WT, cache)
|
||
r.set_worker_mode(WT, "cache")
|
||
startup = r.resolve_startup_config(WT, base, worker_manager=None)
|
||
report(
|
||
"cache 模式：resolve_startup_config 让插件初始化读到缓存值（本次核心修复）",
|
||
startup.get("api_key") == "CACHE-KEY",
|
||
f"api_key={startup.get('api_key')}",
|
||
)
|
||
vs_front = r._read_worker_config_section_with_source(cache, WT)
|
||
report(
|
||
"cache 模式：插件初始化值与前端( get_config )读到的一致",
|
||
vs_front[0].get("api_key") == startup.get("api_key") == "CACHE-KEY"
|
||
and vs_front[1] == "cache",
|
||
f"source={vs_front[1]}",
|
||
)
|
||
|
||
# ── 2. cache 模式：缓存缺失 → 用 toml 引导并写回 ────────────────
|
||
cache2 = {}
|
||
got2 = r.resolve_runtime_config(WT, base, cache2, "cache")
|
||
report(
|
||
"cache 模式：缓存缺失用 toml 引导并写回",
|
||
got2.get("api_key") == "TOML-KEY" and cache2.get(key, {}).get("api_key") == "TOML-KEY",
|
||
f"cache={cache2.get(key)}",
|
||
)
|
||
|
||
# ── 3. toml 模式：以文件为准，启动时覆盖缓存 ────────────────────
|
||
cache3 = {key: dict(cached_cfg)}
|
||
got3 = r.resolve_runtime_config(WT, base, cache3, "toml")
|
||
report(
|
||
"toml 模式：返回值取 toml 值",
|
||
got3.get("api_key") == "TOML-KEY" and got3.get("model") == "toml-model",
|
||
f"api_key={got3.get('api_key')}",
|
||
)
|
||
report(
|
||
"toml 模式：启动时用 toml 覆盖缓存（前端随之显示 toml 值）",
|
||
cache3.get(key, {}).get("api_key") == "TOML-KEY",
|
||
f"cache={cache3.get(key)}",
|
||
)
|
||
r.set_worker_cache(WT, cache3)
|
||
r.set_worker_mode(WT, "toml")
|
||
startup3 = r.resolve_startup_config(WT, base, worker_manager=None)
|
||
report(
|
||
"toml 模式：插件初始化值与缓存/前端一致（均为 toml 值）",
|
||
startup3.get("api_key") == "TOML-KEY"
|
||
and r._read_worker_config_section(cache3, WT).get("api_key") == "TOML-KEY",
|
||
f"api_key={startup3.get('api_key')}",
|
||
)
|
||
|
||
# ── 4. 无缓存：保持 toml 种子 ────────────────────────────────
|
||
got4 = r.resolve_runtime_config(WT, base, None, "toml")
|
||
report(
|
||
"无缓存后端：保持 toml 种子（经典行为）",
|
||
got4.get("api_key") == "TOML-KEY",
|
||
f"api_key={got4.get('api_key')}",
|
||
)
|
||
|
||
# ── 5. get_config 来源标注 ──────────────────────────────────
|
||
only_missing_key = {"worker:cm:configs:other": {"api_key": "X"}}
|
||
cfg5, src5 = r._read_worker_config_section_with_source(only_missing_key, WT)
|
||
report(
|
||
"来源标注：缓存对象存在但键缺失 → 回退 toml 且标 toml",
|
||
src5 == "toml" and cfg5.get("api_key") == "TOML-KEY",
|
||
f"source={src5}",
|
||
)
|
||
cfg6, src6 = r._read_worker_config_section_with_source({key: {"api_key": "C"}}, WT)
|
||
report("来源标注：命中缓存 → cache", src6 == "cache" and cfg6.get("api_key") == "C", f"source={src6}")
|
||
|
||
# ── 6. count：toml 模式覆盖 / cache 模式以缓存为准 ──────────────
|
||
counts_key = r._SCALE_COUNTS_KEY.format(type=WT)
|
||
cache7 = {counts_key: 1}
|
||
n_toml = r.resolve_worker_count(WT, 3, cache7, "toml")
|
||
report(
|
||
"count toml 模式：用 toml 的 count 覆盖缓存",
|
||
n_toml == 3 and cache7.get(counts_key) == 3,
|
||
f"got={n_toml}, cache={cache7.get(counts_key)}",
|
||
)
|
||
cache8 = {counts_key: 1}
|
||
n_cache = r.resolve_worker_count(WT, 3, cache8, "cache")
|
||
report(
|
||
"count cache 模式：缓存有值用缓存（跨重启保持）",
|
||
n_cache == 1 and cache8.get(counts_key) == 1,
|
||
f"got={n_cache}",
|
||
)
|
||
cache9 = {}
|
||
n_seed = r.resolve_worker_count(WT, 3, cache9, "cache")
|
||
report(
|
||
"count cache 模式：缓存缺失用 toml 引导并写回",
|
||
n_seed == 3 and cache9.get(counts_key) == 3,
|
||
f"got={n_seed}, cache={cache9.get(counts_key)}",
|
||
)
|
||
|
||
# ── 7. Plugin.create_workers：初始化（on_init）前即应用模式语义 ──────
|
||
# 这是本次主 bug：cache 模式下插件此前按 toml 建客户端，只有 RPC/管理通道读缓存。
|
||
import asyncio
|
||
|
||
from chongming_worker.plugin import Plugin
|
||
|
||
seen = {}
|
||
|
||
class DemoPlugin(Plugin):
|
||
type = WT
|
||
auto_bind = False
|
||
|
||
async def _bootstrap(self): # 隔离数据库/对象存储引导，只验证配置来源
|
||
return None
|
||
|
||
async def on_init(self):
|
||
seen["on_init"] = self.config.get("api_key")
|
||
|
||
def _create_workers_impl(self, session, config, count=1, worker_manager=None):
|
||
seen["impl"] = (config or {}).get("api_key")
|
||
return []
|
||
|
||
r.set_worker_cache(WT, {key: {"api_key": "CACHE-KEY"}})
|
||
r.set_worker_mode(WT, "cache")
|
||
asyncio.run(DemoPlugin().create_workers(None, dict(base), 1))
|
||
report(
|
||
"Plugin.create_workers：cache 模式 on_init 读到缓存值",
|
||
seen.get("on_init") == "CACHE-KEY",
|
||
f"on_init={seen.get('on_init')}",
|
||
)
|
||
report(
|
||
"Plugin.create_workers：下发给 WorkerPool 的配置同为缓存值",
|
||
seen.get("impl") == "CACHE-KEY",
|
||
f"impl={seen.get('impl')}",
|
||
)
|
||
|
||
r.set_worker_cache(WT, {key: {"api_key": "CACHE-KEY"}})
|
||
r.set_worker_mode(WT, "toml")
|
||
seen.clear()
|
||
asyncio.run(DemoPlugin().create_workers(None, dict(base), 1))
|
||
report(
|
||
"Plugin.create_workers：toml 模式 on_init 读到文件值",
|
||
seen.get("on_init") == "TOML-KEY" and seen.get("impl") == "TOML-KEY",
|
||
f"on_init={seen.get('on_init')}, impl={seen.get('impl')}",
|
||
)
|
||
|
||
print("\n" + "=" * 56)
|
||
print(f"测试结果: ✅ {PASS} 通过 ❌ {FAIL} 失败")
|
||
return 0 if FAIL == 0 else 1
|
||
finally:
|
||
os.chdir(cwd)
|
||
shutil.rmtree(workdir, ignore_errors=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(run())
|