diff --git a/governor.py b/governor.py index 27102d88..7e62ff8b 100755 --- a/governor.py +++ b/governor.py @@ -36,8 +36,8 @@ class Governor: self.postgresql = Postgresql(config['postgresql']) self.ha = Ha(self.postgresql, self.etcd) - def touch_member(self): - return self.etcd.touch_member(self.postgresql.name, self.postgresql.connection_string) + def touch_member(self, ttl=None): + return self.etcd.touch_member(self.postgresql.name, self.postgresql.connection_string, ttl) def initialize(self): # wait for etcd to be available @@ -72,6 +72,10 @@ class Governor: def main(): + logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO) + signal.signal(signal.SIGTERM, sigterm_handler) + signal.signal(signal.SIGCHLD, sigchld_handler) + if len(sys.argv) < 2 or not os.path.isfile(sys.argv[1]): print('Usage: {} config.yml'.format(sys.argv[0])) return @@ -85,12 +89,10 @@ def main(): RestApiServer(governor).start() governor.run() finally: + governor.touch_member(300) # schedule member removal governor.postgresql.stop() governor.etcd.delete_leader(governor.postgresql.name) if __name__ == '__main__': - logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO) - signal.signal(signal.SIGTERM, sigterm_handler) - signal.signal(signal.SIGCHLD, sigchld_handler) main() diff --git a/tests/test_governor.py b/tests/test_governor.py index ed7ae32e..844f374e 100644 --- a/tests/test_governor.py +++ b/tests/test_governor.py @@ -6,7 +6,7 @@ import sys import time import yaml -from governor import Governor, main, sigchld_handler +from governor import Governor, main, sigchld_handler, sigterm_handler from test_ha import true, false from test_postgresql import Postgresql, os_system, psycopg2_connect from test_etcd import requests_get, requests_put, requests_delete @@ -20,6 +20,10 @@ def os_waitpid(a, b): return (0, 0) +def time_sleep(_): + raise Exception() + + class TestGovernor(unittest.TestCase): def __init__(self, method_name='runTest'): @@ -28,13 +32,13 @@ class TestGovernor(unittest.TestCase): super(TestGovernor, self).__init__(method_name) def set_up(self): + self.touched = False os.system = os_system psycopg2.connect = psycopg2_connect requests.get = requests_get requests.put = requests_put requests.delete = requests_delete time.sleep = nop - Governor.run = nop self.write_pg_hba = Postgresql.write_pg_hba self.write_recovery_conf = Postgresql.write_recovery_conf Postgresql.write_pg_hba = nop @@ -44,9 +48,20 @@ class TestGovernor(unittest.TestCase): Postgresql.write_pg_hba = self.write_pg_hba Postgresql.write_recovery_conf = self.write_recovery_conf + def test_sigterm_handler(self): + self.assertRaises(SystemExit, sigterm_handler, None, None) + def test_governor_main(self): - sys.argv = ['governor.py', 'postgres0.yml'] main() + sys.argv = ['governor.py', 'postgres0.yml'] + time.sleep = time_sleep + self.assertRaises(Exception, main) + + def touch_member(self): + if not self.touched: + self.touched = True + return False + return True def test_governor_initialize(self): with open('postgres0.yml', 'r') as f: @@ -60,7 +75,12 @@ class TestGovernor(unittest.TestCase): g.etcd.race = false g.initialize() g.postgresql.data_directory_empty = false + g.touch_member = self.touch_member g.initialize() + g.postgresql.data_directory_empty = true + time.sleep = time_sleep + g.postgresql.sync_from_leader = false + self.assertRaises(Exception, g.initialize) def test_sigchld_handler(self): sigchld_handler(None, None)