200 lines
6.5 KiB
Python
200 lines
6.5 KiB
Python
"""Unit tests for the single-instance leader election (singleton.py)."""
|
|
import fcntl
|
|
import importlib
|
|
import os
|
|
|
|
import psycopg
|
|
|
|
from core import singleton
|
|
|
|
|
|
def _fresh():
|
|
# _held is module-level state; reload for an isolated starting point per test.
|
|
return importlib.reload(singleton)
|
|
|
|
|
|
def test_no_lock_path_is_always_leader():
|
|
s = _fresh()
|
|
assert s.claim(None) is True
|
|
assert s.claim("") is True
|
|
|
|
|
|
def test_first_claimer_wins_and_holds_the_lock(tmp_path):
|
|
s = _fresh()
|
|
lock = str(tmp_path / "notifier.lock")
|
|
assert s.claim(lock) is True
|
|
# Held for the process lifetime: the fd stays open and flock'd.
|
|
assert lock in s._held
|
|
|
|
|
|
def test_reclaiming_same_path_is_idempotent(tmp_path):
|
|
s = _fresh()
|
|
lock = str(tmp_path / "notifier.lock")
|
|
assert s.claim(lock) is True
|
|
fd = s._held[lock]
|
|
# Re-claiming must not re-open or drop the lock — same fd, still True.
|
|
assert s.claim(lock) is True
|
|
assert s._held[lock] == fd
|
|
|
|
|
|
def test_second_holder_stands_down(tmp_path):
|
|
"""When another holder already owns the lock, claim() returns False and leaks no
|
|
fd. Simulate the other worker by taking the exclusive flock on a separate open
|
|
file description (a fresh os.open, as a distinct process would)."""
|
|
s = _fresh()
|
|
lock = str(tmp_path / "notifier.lock")
|
|
other = os.open(lock, os.O_RDWR | os.O_CREAT, 0o644)
|
|
fcntl.flock(other, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
try:
|
|
assert s.claim(lock) is False
|
|
assert lock not in s._held # did not record a role it doesn't own
|
|
finally:
|
|
fcntl.flock(other, fcntl.LOCK_UN)
|
|
os.close(other)
|
|
|
|
|
|
def test_reelection_after_holder_releases(tmp_path):
|
|
"""A restart drops the OS lock; the next claim() then wins."""
|
|
s = _fresh()
|
|
lock = str(tmp_path / "notifier.lock")
|
|
other = os.open(lock, os.O_RDWR | os.O_CREAT, 0o644)
|
|
fcntl.flock(other, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
assert s.claim(lock) is False
|
|
# Prior leader exits -> OS releases the lock.
|
|
fcntl.flock(other, fcntl.LOCK_UN)
|
|
os.close(other)
|
|
assert s.claim(lock) is True
|
|
|
|
|
|
# --- claim_pg (Postgres advisory lock) --------------------------------------
|
|
# Hermetic: no real Postgres. IS_POSTGRES/_PG_DSN are set directly on the reloaded
|
|
# module (mirroring how tests elsewhere monkeypatch module state), and
|
|
# psycopg.connect is stubbed so a "connection" is a bare object carrying only
|
|
# what claim_pg touches: .execute(...).fetchone() and .closed/.close().
|
|
|
|
class _FakeCursor:
|
|
def __init__(self, won):
|
|
self._won = won
|
|
|
|
def fetchone(self):
|
|
return (self._won,)
|
|
|
|
|
|
class _FakeConn:
|
|
def __init__(self, won):
|
|
self.closed = False
|
|
self._won = won
|
|
|
|
def execute(self, *a, **kw):
|
|
return _FakeCursor(self._won)
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
|
|
def _postgres(s, monkeypatch):
|
|
monkeypatch.setattr(s, "IS_POSTGRES", True)
|
|
monkeypatch.setattr(s, "_PG_DSN", "postgresql://fake/dsn")
|
|
return s
|
|
|
|
|
|
def test_claim_pg_false_when_not_postgres(monkeypatch):
|
|
s = _fresh()
|
|
monkeypatch.setattr(s, "IS_POSTGRES", False)
|
|
assert s.claim_pg(1) is False
|
|
assert 1 not in s._pg_held
|
|
|
|
|
|
def test_claim_pg_wins_and_holds_the_connection(monkeypatch):
|
|
s = _postgres(_fresh(), monkeypatch)
|
|
monkeypatch.setattr(psycopg, "connect", lambda *a, **kw: _FakeConn(won=True))
|
|
assert s.claim_pg(1) is True
|
|
assert 1 in s._pg_held and s._pg_held[1].closed is False
|
|
|
|
|
|
def test_claim_pg_reclaiming_is_idempotent_without_reconnecting(monkeypatch):
|
|
s = _postgres(_fresh(), monkeypatch)
|
|
held = _FakeConn(won=True)
|
|
s._pg_held[1] = held
|
|
|
|
def _should_not_be_called(*a, **kw):
|
|
raise AssertionError("already holding the lock; must not reconnect")
|
|
|
|
monkeypatch.setattr(psycopg, "connect", _should_not_be_called)
|
|
assert s.claim_pg(1) is True
|
|
assert s._pg_held[1] is held
|
|
|
|
|
|
def test_claim_pg_loses_closes_and_does_not_hold(monkeypatch):
|
|
s = _postgres(_fresh(), monkeypatch)
|
|
lost = _FakeConn(won=False)
|
|
monkeypatch.setattr(psycopg, "connect", lambda *a, **kw: lost)
|
|
assert s.claim_pg(1) is False
|
|
assert lost.closed is True
|
|
assert 1 not in s._pg_held
|
|
|
|
|
|
def test_claim_pg_reelects_after_a_dead_held_connection(monkeypatch):
|
|
s = _postgres(_fresh(), monkeypatch)
|
|
dead = _FakeConn(won=True)
|
|
dead.closed = True # e.g. a Postgres restart dropped it
|
|
s._pg_held[1] = dead
|
|
fresh_conn = _FakeConn(won=True)
|
|
monkeypatch.setattr(psycopg, "connect", lambda *a, **kw: fresh_conn)
|
|
assert s.claim_pg(1) is True
|
|
assert s._pg_held[1] is fresh_conn
|
|
|
|
|
|
def test_claim_pg_unreachable_postgres_is_not_leader_not_fatal(monkeypatch):
|
|
s = _postgres(_fresh(), monkeypatch)
|
|
|
|
def _boom(*a, **kw):
|
|
raise OSError("connection refused")
|
|
|
|
monkeypatch.setattr(psycopg, "connect", _boom)
|
|
assert s.claim_pg(1) is False
|
|
assert 1 not in s._pg_held
|
|
|
|
|
|
# --- claim_leader (env-selected dispatch) -----------------------------------
|
|
|
|
def test_claim_leader_defaults_to_always_leader(monkeypatch):
|
|
s = _fresh()
|
|
monkeypatch.delenv("THERMOGRAPH_SINGLETON_PG", raising=False)
|
|
monkeypatch.delenv("THERMOGRAPH_SINGLETON_LOCK", raising=False)
|
|
assert s.claim_leader() is True
|
|
|
|
|
|
def test_claim_leader_uses_flock_when_pg_flag_unset(tmp_path, monkeypatch):
|
|
s = _fresh()
|
|
monkeypatch.delenv("THERMOGRAPH_SINGLETON_PG", raising=False)
|
|
lock = str(tmp_path / "notifier.lock")
|
|
monkeypatch.setenv("THERMOGRAPH_SINGLETON_LOCK", lock)
|
|
assert s.claim_leader() is True
|
|
assert lock in s._held
|
|
|
|
|
|
def test_claim_leader_ignores_pg_flag_when_not_on_postgres(tmp_path, monkeypatch):
|
|
# THERMOGRAPH_SINGLETON_PG is set, but the climate record isn't on Postgres
|
|
# (IS_POSTGRES False) -> falls through to the flock, not claim_pg.
|
|
s = _fresh()
|
|
monkeypatch.setattr(s, "IS_POSTGRES", False)
|
|
monkeypatch.setenv("THERMOGRAPH_SINGLETON_PG", "1")
|
|
lock = str(tmp_path / "notifier.lock")
|
|
monkeypatch.setenv("THERMOGRAPH_SINGLETON_LOCK", lock)
|
|
|
|
def _should_not_be_called(key):
|
|
raise AssertionError("must not use the pg path when IS_POSTGRES is False")
|
|
|
|
monkeypatch.setattr(s, "claim_pg", _should_not_be_called)
|
|
assert s.claim_leader() is True
|
|
assert lock in s._held
|
|
|
|
|
|
def test_claim_leader_dispatches_to_claim_pg_when_configured(monkeypatch):
|
|
s = _postgres(_fresh(), monkeypatch)
|
|
monkeypatch.setenv("THERMOGRAPH_SINGLETON_PG", "1")
|
|
calls = []
|
|
monkeypatch.setattr(s, "claim_pg", lambda key: calls.append(key) or True)
|
|
assert s.claim_leader() is True
|
|
assert calls == [s.NOTIFIER_LOCK_KEY]
|