from collections.abc import Callable from threading import Event from typing import Any from zbt.commands.refund_worker_runtime import ( run_parallel_worker, run_resilient_singleton_worker, run_singleton_worker, ) from zbt.infrastructure.redis.worker_lock import RedisWorkerLease class FakeRedis: def __init__(self, clock: Callable[[], float]) -> None: self._clock = clock self._values: dict[str, tuple[str, float]] = {} def set( self, name: str, value: str, *, nx: bool, ex: int, ) -> bool | None: self._discard_expired(name) if nx and name in self._values: return None self._values[name] = (value, self._clock() + ex) return True def eval(self, script: str, numkeys: int, *keys_and_args: Any) -> int: assert numkeys == 1 name = str(keys_and_args[0]) owner_id = str(keys_and_args[1]) self._discard_expired(name) current = self._values.get(name) if current is None or current[0] != owner_id: return 0 if "'DEL'" in script: del self._values[name] return 1 if "'EXPIRE'" in script: ttl_seconds = int(keys_and_args[2]) self._values[name] = (owner_id, self._clock() + ttl_seconds) return 1 raise NotImplementedError def _discard_expired(self, name: str) -> None: current = self._values.get(name) if current is not None and current[1] <= self._clock(): del self._values[name] def test_only_one_refund_worker_can_acquire_the_same_lease() -> None: now = 100.0 redis = FakeRedis(lambda: now) first = RedisWorkerLease( redis, key="ins:s3:test:worker:lock:refund", owner_id="worker-1", ttl_seconds=30, ) second = RedisWorkerLease( redis, key="ins:s3:test:worker:lock:refund", owner_id="worker-2", ttl_seconds=30, ) assert first.acquire() is True assert second.acquire() is False def test_releasing_the_owned_lease_allows_a_new_worker_to_take_over() -> None: now = 100.0 redis = FakeRedis(lambda: now) first = RedisWorkerLease( redis, key="ins:s3:test:worker:lock:refund", owner_id="worker-1", ttl_seconds=30, ) second = RedisWorkerLease( redis, key="ins:s3:test:worker:lock:refund", owner_id="worker-2", ttl_seconds=30, ) assert first.acquire() is True assert first.release() is True assert second.acquire() is True def test_renewing_the_lease_prevents_takeover_until_the_new_expiry() -> None: now = 100.0 redis = FakeRedis(lambda: now) first = RedisWorkerLease( redis, key="ins:s3:test:worker:lock:refund", owner_id="worker-1", ttl_seconds=30, ) second = RedisWorkerLease( redis, key="ins:s3:test:worker:lock:refund", owner_id="worker-2", ttl_seconds=30, ) assert first.acquire() is True now = 125.0 assert first.renew() is True now = 131.0 assert second.acquire() is False now = 156.0 assert second.acquire() is True def test_expired_owner_cannot_release_the_new_owners_lease() -> None: now = 100.0 redis = FakeRedis(lambda: now) first = RedisWorkerLease( redis, key="ins:s3:test:worker:lock:refund", owner_id="worker-1", ttl_seconds=30, ) second = RedisWorkerLease( redis, key="ins:s3:test:worker:lock:refund", owner_id="worker-2", ttl_seconds=30, ) third = RedisWorkerLease( redis, key="ins:s3:test:worker:lock:refund", owner_id="worker-3", ttl_seconds=30, ) assert first.acquire() is True now = 131.0 assert second.acquire() is True assert first.release() is False assert third.acquire() is False class DelayedLease: owner_id = "worker-2" def __init__(self) -> None: self.acquire_count = 0 self.release_count = 0 def acquire(self) -> bool: self.acquire_count += 1 return self.acquire_count >= 3 def renew(self) -> bool: return True def release(self) -> bool: self.release_count += 1 return True def test_duplicate_worker_waits_and_takes_over_without_early_initialization() -> None: stop_event = Event() lease = DelayedLease() initialized_after_attempt = 0 output: list[str] = [] def create_worker() -> Any: nonlocal initialized_after_attempt initialized_after_attempt = lease.acquire_count class OneShotWorker: def run_once(self, limit: int = 20) -> dict[str, Any]: del limit stop_event.set() return {"processed": 0, "items": []} return OneShotWorker() result = run_singleton_worker( create_worker, lease, stop_event=stop_event, output=output.append, poll_interval_seconds=0.001, heartbeat_interval_seconds=0.001, acquire_retry_interval_seconds=0.001, ) assert result == "stopped" assert lease.acquire_count == 3 assert initialized_after_attempt == 3 assert lease.release_count == 1 assert "已有退款工作器持有单实例锁,当前实例进入待命并等待接管。" in output assert "原实例租约已释放,待命工作器开始接管:worker-2" in output def test_worker_keeps_waiting_when_redis_temporarily_rejects_lease_acquisition() -> None: stop_event = Event() output: list[str] = [] class TemporarilyUnavailableLease: owner_id = "worker-recovering" def __init__(self) -> None: self.acquire_count = 0 def acquire(self) -> bool: self.acquire_count += 1 if self.acquire_count == 1: raise ConnectionError("Redis connection reset") return True def renew(self) -> bool: return True def release(self) -> bool: return True lease = TemporarilyUnavailableLease() class OneShotWorker: def run_once(self, limit: int = 20) -> dict[str, Any]: del limit stop_event.set() return {"processed": 0, "items": []} result = run_singleton_worker( OneShotWorker, lease, stop_event=stop_event, output=output.append, poll_interval_seconds=0.001, heartbeat_interval_seconds=0.01, acquire_retry_interval_seconds=0.001, ) assert result == "stopped" assert lease.acquire_count == 2 assert ( "退款工作器获取单实例锁异常,将继续等待:Redis connection reset" in output ) def test_standby_worker_can_stop_without_initializing_business_dependencies() -> None: stop_event = Event() initialized = False class RejectedLease: owner_id = "worker-standby" def acquire(self) -> bool: stop_event.set() return False def renew(self) -> bool: raise AssertionError("未持锁的工作器不应续约") def release(self) -> bool: raise AssertionError("未持锁的工作器不应释放租约") def create_worker() -> Any: nonlocal initialized initialized = True raise AssertionError("待命工作器不应初始化业务依赖") result = run_singleton_worker( create_worker, RejectedLease(), stop_event=stop_event, output=lambda _: None, acquire_retry_interval_seconds=0.001, ) assert result == "stopped" assert initialized is False class HeldLease: owner_id = "worker-1" def __init__(self) -> None: self.release_count = 0 def acquire(self) -> bool: return True def renew(self) -> bool: return True def release(self) -> bool: self.release_count += 1 return True def test_lock_owner_processes_work_and_releases_the_lease_on_stop() -> None: stop_event = Event() lease = HeldLease() processed = 0 output: list[str] = [] class OneShotWorker: def run_once(self, limit: int = 20) -> dict[str, Any]: nonlocal processed del limit processed += 1 stop_event.set() return {"processed": 1, "items": []} result = run_singleton_worker( OneShotWorker, lease, stop_event=stop_event, output=output.append, poll_interval_seconds=0.001, heartbeat_interval_seconds=0.001, ) assert result == "stopped" assert processed == 1 assert lease.release_count == 1 assert "退款工作器初始化完成,开始轮询任务。" in output def test_active_worker_renews_the_lease_while_business_processing_is_running() -> None: stop_event = Event() renewed = Event() class HeartbeatLease(HeldLease): def renew(self) -> bool: renewed.set() return True class BlockingWorker: def run_once(self, limit: int = 20) -> dict[str, Any]: del limit assert renewed.wait(0.5), "业务执行期间没有续约单实例锁" stop_event.set() return {"processed": 0, "items": []} result = run_singleton_worker( BlockingWorker, HeartbeatLease(), stop_event=stop_event, output=lambda _: None, poll_interval_seconds=0.001, heartbeat_interval_seconds=0.01, ) assert result == "stopped" def test_worker_stops_claiming_tasks_after_losing_the_singleton_lease() -> None: stop_event = Event() renewal_attempted = Event() class LostLease(HeldLease): def renew(self) -> bool: renewal_attempted.set() return False class WorkerFinishingCurrentBatch: def run_once(self, limit: int = 20) -> dict[str, Any]: del limit assert renewal_attempted.wait(0.5) return {"processed": 0, "items": []} result = run_singleton_worker( WorkerFinishingCurrentBatch, LostLease(), stop_event=stop_event, output=lambda _: None, poll_interval_seconds=0.001, heartbeat_interval_seconds=0.01, ) assert result == "lease_lost" def test_resilient_worker_waits_for_the_lease_and_resumes_after_losing_it() -> None: stop_event = Event() first_renewal = Event() initialized = 0 lease_lost_count = 0 output: list[str] = [] class RecoverableLease(HeldLease): def __init__(self) -> None: super().__init__() self.acquire_count = 0 self.renew_count = 0 def acquire(self) -> bool: self.acquire_count += 1 return True def renew(self) -> bool: self.renew_count += 1 if self.renew_count == 1: first_renewal.set() return False return True lease = RecoverableLease() def create_worker() -> Any: nonlocal initialized initialized += 1 current_generation = initialized class RecoveringWorker: def run_once(self, limit: int = 20) -> dict[str, Any]: del limit if current_generation == 1: assert first_renewal.wait(0.5) else: stop_event.set() return {"processed": 0, "items": []} return RecoveringWorker() def record_lease_loss() -> None: nonlocal lease_lost_count lease_lost_count += 1 result = run_resilient_singleton_worker( create_worker, lease, stop_event=stop_event, on_lease_lost=record_lease_loss, output=output.append, poll_interval_seconds=0.001, heartbeat_interval_seconds=0.01, acquire_retry_interval_seconds=0.001, ) assert result == "stopped" assert initialized == 2 assert lease.acquire_count == 2 assert lease_lost_count == 1 assert "退款工作器进入待命,等待重新获取单实例锁。" in output def test_resilient_worker_recovers_when_releasing_a_lost_lease_also_fails() -> None: stop_event = Event() renewal_attempted = Event() initialized = 0 output: list[str] = [] class UnstableRedisLease: owner_id = "worker-unstable-redis" def __init__(self) -> None: self.renew_count = 0 self.release_count = 0 def acquire(self) -> bool: return True def renew(self) -> bool: self.renew_count += 1 if self.renew_count == 1: renewal_attempted.set() raise ConnectionError("Redis unavailable during renew") return True def release(self) -> bool: self.release_count += 1 if self.release_count == 1: raise ConnectionError("Redis unavailable during release") return True lease = UnstableRedisLease() def create_worker() -> Any: nonlocal initialized initialized += 1 current_generation = initialized class RecoveringWorker: def run_once(self, limit: int = 20) -> dict[str, Any]: del limit if current_generation == 1: assert renewal_attempted.wait(0.5) else: stop_event.set() return {"processed": 0, "items": []} return RecoveringWorker() result = run_resilient_singleton_worker( create_worker, lease, stop_event=stop_event, output=output.append, poll_interval_seconds=0.001, heartbeat_interval_seconds=0.01, acquire_retry_interval_seconds=0.001, ) assert result == "stopped" assert initialized == 2 assert lease.release_count == 2 assert ( "退款工作器释放单实例锁异常,租约将由 Redis 自动过期:" "Redis unavailable during release" in output ) def test_lease_is_released_when_worker_initialization_fails() -> None: lease = HeldLease() def fail_to_create_worker() -> Any: raise RuntimeError("数据库初始化失败") try: run_singleton_worker( fail_to_create_worker, lease, stop_event=Event(), output=lambda _: None, ) except RuntimeError as error: assert str(error) == "数据库初始化失败" else: raise AssertionError("工作器初始化异常应继续向上抛出") assert lease.release_count == 1 def test_parallel_mode_processes_tasks_without_a_global_singleton_lease() -> None: stop_event = Event() processed = 0 class OneShotWorker: def run_once(self, limit: int = 20) -> dict[str, Any]: nonlocal processed del limit processed += 1 stop_event.set() return {"processed": 1, "items": []} result = run_parallel_worker( OneShotWorker, stop_event=stop_event, output=lambda _: None, poll_interval_seconds=0.001, ) assert result == "stopped" assert processed == 1