harness_zbt/backend/tests/test_refund_worker_singleton.py

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