thermograph/tests/core/test_singleton.py

"""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]