truf-server/tests/test_observer_only_coordinated_shutdown.py

from concurrent.futures import Future
from pathlib import Path
import sys
import unittest
from unittest import mock


ROOT = Path(__file__).resolve().parents[1]
APP_DIR = ROOT / 'app'
sys.path.insert(0, str(APP_DIR))

import lifecycle_authority
import postgres_runtime
import supervisor


class ImmediateExecutor:
    def submit(self, function):
        future = Future()
        try:
            future.set_result(function())
        except BaseException as exc:
            future.set_exception(exc)
        return future


class LifecycleBackend:
    def __init__(self, probes, starts=()):
        self.probes = list(probes)
        self.starts = list(starts)
        self.start_calls = 0
        self.stop_calls = 0
        self.closed = False

    def probe(self):
        return self.probes.pop(0)

    def start(self):
        self.start_calls += 1
        return self.starts.pop(0)

    def stop(self):
        self.stop_calls += 1
        return postgres_runtime.StopResult(True, True, 'identity-safe backend stop completed')

    def close(self):
        self.closed = True


def ready_controller(owned):
    if owned:
        backend = LifecycleBackend(
            probes=[
                postgres_runtime.ProbeResult(postgres_runtime.ProbeKind.STOPPED, 'offline'),
                postgres_runtime.ProbeResult(postgres_runtime.ProbeKind.READY, 'owned ready', 'owned-epoch'),
            ],
            starts=[postgres_runtime.StartResult(True, 'start accepted')],
        )
    else:
        backend = LifecycleBackend(probes=[
            postgres_runtime.ProbeResult(
                postgres_runtime.ProbeKind.READY,
                'preexisting cluster verified',
                'observer-epoch',
            ),
        ])
    controller = postgres_runtime.PostgresController(
        backend,
        executor=ImmediateExecutor(),
        stable_ready_interval_sec=0,
    )
    for _ in range(4 if owned else 2):
        controller.tick(0)
    assert controller.state == postgres_runtime.PostgresState.READY
    return controller, backend


def shutdown_context(controller):
    return {
        'lifecycle_phase': lifecycle_authority.PHASE_ACTIVE,
        'activation_state': lifecycle_authority.PHASE_ACTIVE,
        'start_gate_open': True,
        'postgres_controller': controller,
        'supervisor_config': {},
    }


class ObserverOnlyCoordinatedShutdownTests(unittest.TestCase):
    def test_preexisting_ready_observer_closes_without_stop(self):
        controller, backend = ready_controller(owned=False)
        source = mock.Mock(source='child')
        source.stop.return_value = True
        source.is_running.return_value = False
        self.assertFalse(controller.lifecycle_action_required)

        with mock.patch.object(controller, 'request_stop', wraps=controller.request_stop) as request_stop:
            self.assertTrue(supervisor.coordinated_shutdown([source], shutdown_context(controller)))

        request_stop.assert_not_called()
        self.assertEqual(backend.stop_calls, 0)
        self.assertTrue(backend.closed)
        self.assertTrue(controller.terminal)
        self.assertTrue(controller.authority_release_safe)
        self.assertIn('observer-only', controller.detail)

    def test_owned_ready_cluster_stops_exactly_once(self):
        controller, backend = ready_controller(owned=True)
        source = mock.Mock(source='child')
        source.stop.return_value = True
        source.is_running.return_value = False
        self.assertTrue(controller.lifecycle_action_required)

        with mock.patch.object(controller, 'request_stop', wraps=controller.request_stop) as request_stop:
            self.assertTrue(supervisor.coordinated_shutdown([source], shutdown_context(controller)))

        request_stop.assert_called_once_with()
        self.assertEqual(backend.stop_calls, 1)
        self.assertTrue(backend.closed)
        self.assertTrue(controller.authority_release_safe)

    def test_child_failure_defers_owned_and_observer_database_action(self):
        for owned in (False, True):
            with self.subTest(owned=owned):
                controller, backend = ready_controller(owned=owned)
                source = mock.Mock(source='child')
                source.stop.return_value = False
                source.is_running.return_value = True
                context = shutdown_context(controller)

                with mock.patch.object(controller, 'request_stop', wraps=controller.request_stop) as request_stop, \
                        mock.patch.object(controller, 'close', wraps=controller.close) as close:
                    self.assertFalse(supervisor.coordinated_shutdown([source], context))

                request_stop.assert_not_called()
                close.assert_not_called()
                self.assertEqual(backend.stop_calls, 0)
                self.assertFalse(backend.closed)
                self.assertEqual(context['lifecycle_phase'], lifecycle_authority.PHASE_FAILED_HOLD)
                self.assertFalse(context['authority_release_safe'])

    def test_failed_hold_retry_converges_by_ownership(self):
        for owned, expected_stops in ((False, 0), (True, 1)):
            with self.subTest(owned=owned):
                controller, backend = ready_controller(owned=owned)
                source = mock.Mock(source='child')
                source.stop.side_effect = [False, True]
                source.is_running.side_effect = [True, False]
                context = shutdown_context(controller)

                self.assertFalse(supervisor.coordinated_shutdown([source], context))
                self.assertTrue(supervisor.retain_unsafe_authority([source], context, max_attempts=1))

                self.assertEqual(backend.stop_calls, expected_stops)
                self.assertTrue(backend.closed)
                self.assertTrue(controller.terminal)
                self.assertTrue(controller.authority_release_safe)
                self.assertTrue(context['authority_release_safe'])


if __name__ == '__main__':
    unittest.main()