Improve test coverage for governor.py

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()

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)