80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
"""
|
||
Redis 协调骨架自检（默认 Memory 后端）。
|
||
|
||
用法：
|
||
set REDIS_BACKEND=memory
|
||
python -m agent.redis_coord.selftest
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import os
|
||
import sys
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def main() -> int:
|
||
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
|
||
os.environ["REDIS_BACKEND"] = "memory"
|
||
os.environ.setdefault("YTD_ENV", "test")
|
||
os.environ.setdefault("REDIS_KEY_PREFIX", "inquiry_robot:")
|
||
|
||
from agent.config import get_settings
|
||
from agent.redis_coord import (
|
||
get_redis_runtime,
|
||
get_wecom_token,
|
||
rate_limit_allow,
|
||
set_wecom_token,
|
||
shutdown_redis_runtime,
|
||
try_acquire_idempotency,
|
||
)
|
||
from agent.redis_coord.keys import join_key
|
||
|
||
get_settings.cache_clear()
|
||
# 证明：名字在前缀，不在「库名」
|
||
k = join_key("inquiry_robot:", "stream", "llm")
|
||
assert k.startswith("inquiry_robot:"), k
|
||
logger.info("OK: 键前缀示例 %s（db 号在 REDIS_URL 的 /1，不在键名里）", k)
|
||
|
||
rt = get_redis_runtime(get_settings())
|
||
try:
|
||
assert rt.client.ping()
|
||
logger.info("OK: ping backend=%s", rt.client.backend)
|
||
|
||
assert try_acquire_idempotency(
|
||
rt.client, scope="selftest", idempotency_key="once", ttl_seconds=60
|
||
)
|
||
assert not try_acquire_idempotency(
|
||
rt.client, scope="selftest", idempotency_key="once", ttl_seconds=60
|
||
)
|
||
logger.info("OK: 短幂等")
|
||
|
||
r1 = rate_limit_allow(rt.client, scope="selftest", limit=2, window_seconds=60)
|
||
r2 = rate_limit_allow(rt.client, scope="selftest", limit=2, window_seconds=60)
|
||
r3 = rate_limit_allow(rt.client, scope="selftest", limit=2, window_seconds=60)
|
||
assert r1.allowed and r2.allowed and not r3.allowed
|
||
logger.info("OK: 限流")
|
||
|
||
set_wecom_token(rt.client, agent_id="1000010", token="tok-selftest", ttl_seconds=30)
|
||
assert get_wecom_token(rt.client, agent_id="1000010") == "tok-selftest"
|
||
logger.info("OK: token 缓存")
|
||
|
||
eid = rt.llm_stream.enqueue(
|
||
payload={"mode": "ordered_actions", "text": "hi"},
|
||
idempotency_key="msg-1",
|
||
)
|
||
claimed = rt.llm_stream.claim(consumer="selftest-worker", count=1, block_ms=10)
|
||
assert claimed and claimed[0][0] == eid
|
||
rt.llm_stream.ack(eid)
|
||
logger.info("OK: Stream 入队/认领/ack id=%s", eid)
|
||
logger.info("Redis 骨架自检通过")
|
||
return 0
|
||
finally:
|
||
shutdown_redis_runtime()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|