diff --git a/src/resilient_http_client/failure_store.py b/src/resilient_http_client/failure_store.py index a332a89..a83d737 100644 --- a/src/resilient_http_client/failure_store.py +++ b/src/resilient_http_client/failure_store.py @@ -38,13 +38,13 @@ async def set_state(self, state: str, ttl: int | None = None): key_state = self._key("state") key_cd = self._key("open_cooldown") - if state == CircuitState.OPEN.value and ttl: - # issue: model is not atomic - what if this block is interrupted? - await self.redis.set(key_state, CircuitState.OPEN.value) - await self.redis.set(key_cd, "1", ex=ttl, nx=True) - else: - await self.redis.set(key_state, state) - await self.redis.delete(key_cd) + async with self.redis.pipeline(transaction=True) as pipe: + pipe.set(key_state, state) + if state == CircuitState.OPEN.value and ttl: + pipe.set(key_cd, "1", ex=ttl, nx=True) + else: + pipe.delete(key_cd) + await pipe.execute() except Exception as e: logger.error(f"Redis error in set_state: {e}") diff --git a/tests/test_failure_store.py b/tests/test_failure_store.py index f3bcb99..95bf75b 100644 --- a/tests/test_failure_store.py +++ b/tests/test_failure_store.py @@ -3,10 +3,41 @@ from resilient_http_client import CircuitState, FailureStore +class FakePipeline: + def __init__(self, redis): + self.redis = redis + self.commands = [] + + def set(self, key, value, ex=None, nx=False): + self.commands.append(("set", (key, value), {"ex": ex, "nx": nx})) + return self + + def delete(self, key): + self.commands.append(("delete", (key,), {})) + return self + + async def execute(self): + results = [] + for cmd, args, kwargs in self.commands: + method = getattr(self.redis, cmd) + res = await method(*args, **kwargs) + results.append(res) + return results + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + pass + + class FakeRedis: def __init__(self): self.db = {} + def pipeline(self, transaction=True): + return FakePipeline(self) + async def get(self, key): return self.db.get(key) @@ -93,3 +124,17 @@ async def test_half_open_metrics(): await store.reset_half_open() assert await store.get_half_open_calls() == 0 assert await store.get_half_open_successes() == 0 + + +@pytest.mark.asyncio +async def test_atomic_state_transition(): + redis = FakeRedis() + store = FailureStore(redis, "stripe") + + await store.set_state(CircuitState.OPEN.value, ttl=60) + assert await store.get_state() == CircuitState.OPEN.value + assert await redis.get("resilience:stripe:open_cooldown") == "1" + + await store.set_state(CircuitState.CLOSED.value) + assert await store.get_state() == CircuitState.CLOSED.value + assert await redis.get("resilience:stripe:open_cooldown") is None diff --git a/tests/test_flaky_resilience.py b/tests/test_flaky_resilience.py index e2481f1..cc8a148 100644 --- a/tests/test_flaky_resilience.py +++ b/tests/test_flaky_resilience.py @@ -12,10 +12,41 @@ # ----------------------------- # Fake Redis (isolated unit test store) # ----------------------------- +class FakePipeline: + def __init__(self, redis): + self.redis = redis + self.commands = [] + + def set(self, key, value, ex=None, nx=False): + self.commands.append(("set", (key, value), {"ex": ex, "nx": nx})) + return self + + def delete(self, key): + self.commands.append(("delete", (key,), {})) + return self + + async def execute(self): + results = [] + for cmd, args, kwargs in self.commands: + method = getattr(self.redis, cmd) + res = await method(*args, **kwargs) + results.append(res) + return results + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + pass + + class FakeRedis: def __init__(self): self.db = {} + def pipeline(self, transaction=True): + return FakePipeline(self) + async def get(self, key): return self.db.get(key) diff --git a/tests/test_integration.py b/tests/test_integration.py index dfb3428..d74970a 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -8,10 +8,41 @@ ) +class FakePipeline: + def __init__(self, redis): + self.redis = redis + self.commands = [] + + def set(self, key, value, ex=None, nx=False): + self.commands.append(("set", (key, value), {"ex": ex, "nx": nx})) + return self + + def delete(self, key): + self.commands.append(("delete", (key,), {})) + return self + + async def execute(self): + results = [] + for cmd, args, kwargs in self.commands: + method = getattr(self.redis, cmd) + res = await method(*args, **kwargs) + results.append(res) + return results + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + pass + + class FakeRedis: def __init__(self): self.db = {} + def pipeline(self, transaction=True): + return FakePipeline(self) + async def get(self, key): return self.db.get(key)