Merge branch 'restapi' of github.com:CyberDem0n/governor into features/refactoring

Conflicts:
	governor.py
	helpers/postgresql.py
	postgres0.yml
	postgres1.yml
	tests/test_postgresql.py

governor.py

@@ -4,14 +4,13 @@ import logging
 import os
 import signal
 import sys
-import threading
 import time
 import yaml
 
+from helpers.api import RestApiServer
 from helpers.etcd import Etcd
 from helpers.postgresql import Postgresql
 from helpers.ha import Ha
-from helpers.statuspage import getHTTPServer
 
 
 def sigterm_handler(signo, stack_frame):
@@ -37,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
@@ -74,6 +73,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
@@ -82,25 +85,17 @@ def main():
         config = yaml.load(f)
 
     governor = Governor(config)
-
-    # Start the http_server to serve a simple healthcheck
-    http_server = getHTTPServer(governor.postgresql, http_port=config.get(
-        'healtcheck_port', 8008), listen_address='0.0.0.0')
-    http_thread = threading.Thread(target=http_server.serve_forever, args=())
-    http_thread.daemon = True
-
-    governor.initialize()
-    http_thread.start()
-
     try:
+        governor.initialize()
+        # Start the http_server to serve a simple healthcheck
+        host, port = config['restapi']['listen'].split(':')
+        RestApiServer(governor, host, int(port)).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()

helpers/api.py

@@ -0,0 +1,70 @@
+import json
+import logging
+import psycopg2
+import sys
+
+from threading import Thread
+
+if sys.hexversion >= 0x03000000:
+    from http.server import BaseHTTPRequestHandler, HTTPServer
+else:
+    from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
+
+
+logger = logging.getLogger(__name__)
+
+
+class RestApiHandler(BaseHTTPRequestHandler):
+
+    def do_GET(self):
+        try:
+            response = self.get_postgresql_status()
+        except (psycopg2.OperationalError, psycopg2.InterfaceError):
+            logging.exception('get_postgresql_status')
+            response = {'running': False}
+
+        path = '/master' if self.path == '/' else self.path
+        status_code = 200 if response['running'] and response['role'] in path else 503
+
+        self.send_response(status_code)
+        self.send_header('Content-Type', 'application/json')
+        self.end_headers()
+        self.wfile.write(json.dumps(response).encode('utf-8'))
+
+    def get_postgresql_status(self):
+        if not self.server.governor.postgresql.is_running():
+            return {'running': False}
+        cursor = self.server.cursor()
+        cursor.execute("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
+                                 pg_is_in_recovery(),
+                                 pg_current_xlog_location(),
+                                 pg_last_xlog_receive_location(),
+                                 pg_last_xlog_replay_location(),
+                                 pg_is_in_recovery() AND pg_is_xlog_replay_paused()""")
+        row = cursor.fetchone()
+        return {
+            'running': True,
+            'postmaster_start_time': row[0],
+            'role': 'slave' if row[1] else 'master',
+            'xlog': ({
+                'received_location': row[3],
+                'replayed_location': row[4],
+                'paused': row[5]} if row[1] else {
+                'location': row[2]
+            })
+        }
+
+
+class RestApiServer(HTTPServer, Thread):
+
+    def __init__(self, governor, listen_address='0.0.0.0', listen_port=8080):
+        HTTPServer.__init__(self, (listen_address, listen_port), RestApiHandler)
+        Thread.__init__(self, target=self.serve_forever)
+        self.governor = governor
+        self._cursor_holder = None
+        self.daemon = True
+
+    def cursor(self):
+        if not self._cursor_holder or self._cursor_holder.closed:
+            self._cursor_holder = self.governor.postgresql.connection().cursor()
+        return self._cursor_holder

helpers/errors.py

@@ -9,7 +9,3 @@ class EtcdError(Exception):
 
 class CurrentLeaderError(EtcdError):
     pass
-
-
-class HealthiestMemberError(EtcdError):
-    pass

helpers/etcd.py

@@ -2,16 +2,16 @@ import logging
 import requests
 import time
 
+from requests.exceptions import RequestException
 from collections import namedtuple
 from helpers.errors import CurrentLeaderError, EtcdError
 
 logger = logging.getLogger(__name__)
 
-
 Member = namedtuple('Member', 'hostname,address,ttl')
 
 
-class Cluster(namedtuple('Cluster', 'leader,last_leader_operation,members')):
+class Cluster(namedtuple('Cluster', 'initialize,leader,last_leader_operation,members')):
 
     def is_unlocked(self):
         return not (self.leader and self.leader.hostname)
@@ -35,7 +35,7 @@ class Etcd:
                 response = requests.get(self.client_url(path))
                 if response.status_code == 200:
                     break
-            except Exception as e:
+            except RequestException as e:
                 logger.exception('get_client_path')
                 ex = e
 
@@ -54,15 +54,15 @@ class Etcd:
         try:
             response = requests.put(self.client_url(path), data=data)
             return response.status_code in [200, 201, 202, 204]
-        except:
+        except RequestException:
             logger.exception('PUT %s data=%s', path, data)
-            return False
+        raise EtcdError('Etcd is not responding properly')
 
     def delete_client_path(self, path):
         try:
             response = requests.delete(self.client_url(path))
             return response.status_code in [200, 202, 204]
-        except:
+        except RequestException:
             logger.exception('DELETE %s', path)
             return False
 
@@ -87,6 +87,8 @@ class Etcd:
         try:
             response, status_code = self.get_client_path('?recursive=true')
             if status_code == 200:
+                node = self.find_node(response['node'], '/initialize')
+                initialize = True if node else False
                 # get list of members
                 node = self.find_node(response['node'], '/members') or {'nodes': []}
                 members = [Member(n['key'].split('/')[-1], n['value'], n.get('ttl', None)) for n in node['nodes']]
@@ -108,11 +110,11 @@ class Etcd:
                             leader = m
                             break
                     if not leader:
-                        leader = Member(leader['value'], None, None)
+                        leader = Member(node['value'], None, None)
 
-                return Cluster(leader, last_leader_operation, members)
+                return Cluster(initialize, leader, last_leader_operation, members)
             elif status_code == 404:
-                return Cluster(None, None, [])
+                return Cluster(False, None, None, [])
         except:
             logger.exception('get_cluster')
 
@@ -122,27 +124,43 @@ class Etcd:
         try:
             cluster = self.get_cluster()
             return None if cluster.is_unlocked() else cluster.leader
-        except:
-            raise CurrentLeaderError("Etcd is not responding properly")
+        except EtcdError:
+            raise CurrentLeaderError('Etcd is not responding properly')
 
-    def touch_member(self, member, connection_string):
-        return self.put_client_path('/members/' + member, value=connection_string, ttl=self.member_ttl)
+    def touch_member(self, member, connection_string, ttl=None):
+        try:
+            return self.put_client_path('/members/' + member, value=connection_string, ttl=ttl or self.member_ttl)
+        except EtcdError:
+            return False
 
     def take_leader(self, value):
+        try:
             return self.put_client_path('/leader', value=value, ttl=self.ttl)
+        except EtcdError:
+            return False
 
     def attempt_to_acquire_leader(self, value):
+        try:
             ret = self.put_client_path('/leader', value=value, ttl=self.ttl, prevExist=False)
             ret or logger.info('Could not take out TTL lock')
             return ret
+        except EtcdError:
+            return False
 
     def update_leader(self, state_handler):
-        ret = self.put_client_path('/leader', value=state_handler.name, ttl=self.ttl, prevValue=state_handler.name)
-        ret and self.put_client_path('/optime/leader', value=state_handler.last_operation())
-        return ret
+        if self.put_client_path('/leader', value=state_handler.name, ttl=self.ttl, prevValue=state_handler.name):
+            try:
+                self.put_client_path('/optime/leader', value=state_handler.last_operation())
+            except EtcdError:
+                pass
+            return True
+        return False
 
     def race(self, path, value):
+        try:
             return self.put_client_path(path, value=value, prevExist=False)
+        except EtcdError:
+            return False
 
     def delete_member(self, member):
         return self.delete_client_path('/members/' + member)

helpers/ha.py

@@ -1,7 +1,7 @@
 import logging
 
-from helpers.errors import EtcdError, HealthiestMemberError
-from psycopg2 import OperationalError
+from helpers.errors import EtcdError
+from psycopg2 import InterfaceError, OperationalError
 
 logger = logging.getLogger(__name__)
 
@@ -92,7 +92,5 @@ class Ha:
             if self.state_handler.is_leader():
                 self.state_handler.demote(None)
                 return 'demoted self because etcd is not accessible and i was a leader'
-        except OperationalError:
+        except (InterfaceError, OperationalError):
             logger.error('Error communicating with Postgresql.  Will try again')
-        except HealthiestMemberError:
-            logger.error('failed to determine healthiest member fromt etcd')

helpers/postgresql.py

@@ -19,26 +19,27 @@ logger = logging.getLogger(__name__)
 
 def parseurl(url):
     r = urlparse(url)
-    return {
-        'hostname': r.hostname,
+    ret = {
+        'host': r.hostname,
         'port': r.port or 5432,
-        'username': r.username,
-        'password': r.password,
+        'database': r.path[1:],
+        'fallback_application_name': 'Governor',
+        'connect_timeout': 3,
+        'options': '-c statement_timeout=2000',
     }
+    if r.username:
+        ret['user'] = r.username
+    if r.password:
+        ret['password'] = r.password
+    return ret
 
 
 class Postgresql:
 
     def __init__(self, config):
+        self.config = config
         self.name = config['name']
         self.listen_addresses, self.port = config['listen'].split(':')
-        self.libpq_parameters = {
-            'host': self.listen_addresses.split(',')[0].strip(),
-            'port': self.port,
-            'fallback_application_name': 'Governor',
-            'connect_timeout': 5,
-            'options': '-c statement_timeout=2000'
-        }
         self.data_dir = config['data_dir']
         self.replication = config['replication']
         self.superuser = config['superuser']
@@ -46,51 +47,62 @@ class Postgresql:
         self.recovery_conf = os.path.join(self.data_dir, 'recovery.conf')
         self.configuration_to_save = (os.path.join(self.data_dir, 'pg_hba.conf'),
                                       os.path.join(self.data_dir, 'postgresql.conf'))
+        self.pid_path = os.path.join(self.data_dir, 'postmaster.pid')
         self._pg_ctl = 'pg_ctl -w -D ' + self.data_dir
         self.wal_e = config.get('wal_e', None)
         if self.wal_e:
             self.wal_e_path = 'envdir {} wal-e --aws-instance-profile '.\
                 format(self.wal_e.get('env_dir', '/home/postgres/etc/wal-e.d/env'))
 
-        self.config = config
-
-        connect_address = config.get('connect_address', config['listen']) or config['listen']
+        self.local_address = self.get_local_address()
+        connect_address = config.get('connect_address', None) or self.local_address
         self.connection_string = 'postgres://{username}:{password}@{connect_address}/postgres'.format(
             connect_address=connect_address, **self.replication)
 
-        self.conn = None
-        self.cursor_holder = None
+        self._connection = None
+        self._cursor_holder = None
         self.members = []  # list of already existing replication slots
 
-    def cursor(self):
-        if not self.cursor_holder:
-            self.conn = psycopg2.connect(**self.libpq_parameters)
-            self.conn.autocommit = True
-            self.cursor_holder = self.conn.cursor()
+    def get_local_address(self):
+        # TODO: try to get unix_socket_directory from postmaster.pid
+        return self.listen_addresses.split(',')[0].strip() + ':' + self.port
 
-        return self.cursor_holder
+    def connection(self):
+        if not self._connection or self._connection.closed != 0:
+            r = parseurl('postgres://{}/postgres'.format(self.local_address))
+            self._connection = psycopg2.connect(**r)
+            self._connection.autocommit = True
+        return self._connection
+
+    def _cursor(self):
+        if not self._cursor_holder or self._cursor_holder.closed:
+            self._cursor_holder = self.connection().cursor()
+        return self._cursor_holder
 
     def disconnect(self):
-        try:
-            self.conn.close()
-        except:
-            logger.exception('Error disconnecting')
+        self._connection and self._connection.close()
+        self._connection = self._cursor_holder = None
 
     def query(self, sql, *params):
         max_attempts = 0
         while True:
+            ex = None
             try:
-                self.cursor().execute(sql, params)
-                break
+                cursor = self._cursor()
+                cursor.execute(sql, params)
+                return cursor
+            except psycopg2.InterfaceError as e:
+                ex = e
             except psycopg2.OperationalError as e:
-                if self.conn:
-                    self.disconnect()
-                self.cursor_holder = None
-                if max_attempts > 4:
+                if self._connection and self._connection.closed == 0:
                     raise e
+                ex = e
+            if ex:
+                self.disconnect()
                 max_attempts += 1
+                if max_attempts >= 3:
+                    raise ex
                 time.sleep(5)
-        return self.cursor()
 
     def data_directory_empty(self):
         return not os.path.exists(self.data_dir) or os.listdir(self.data_dir) == []
@@ -106,17 +118,17 @@ class Postgresql:
         pgpass = 'pgpass'
         with open(pgpass, 'w') as f:
             os.fchmod(f.fileno(), 0o600)
-            f.write('{hostname}:{port}:*:{username}:{password}\n'.format(**r))
+            f.write('{host}:{port}:*:{user}:{password}\n'.format(**r))
 
         try:
             os.environ['PGPASSFILE'] = pgpass
-            return self.create_replica(leader.address, r) == 0
+            return self.create_replica(r) == 0
         finally:
             os.environ.pop('PGPASSFILE')
 
-    def create_replica(self, master_connurl, master_connection):
+    def create_replica(self, master_connection):
         """ creates a new replica using either pg_basebackup or WAL-E """
-        if self.should_use_s3_to_create_replica(master_connurl):
+        if self.should_use_s3_to_create_replica(master_connection):
             result = self.create_replica_with_s3()
             # if restore from the backup on S3 failed - try with the pg_basebackup
             if result == 0:
@@ -124,7 +136,7 @@ class Postgresql:
         return self.create_replica_with_pg_basebackup(master_connection)
 
     def create_replica_with_pg_basebackup(self, master_connection):
-        return os.system('pg_basebackup -R -D {data_dir} --host={hostname} --port={port} -U {username}'.format(
+        return os.system('pg_basebackup -R -D {data_dir} --host={host} --port={port} -U {user}'.format(
             data_dir=self.data_dir, **master_connection))
 
     def create_replica_with_s3(self):
@@ -135,7 +147,7 @@ class Postgresql:
         self.restore_configuration_files()
         return ret
 
-    def should_use_s3_to_create_replica(self, master_connurl):
+    def should_use_s3_to_create_replica(self, master_connection):
         """ determine whether it makes sense to use S3 and not pg_basebackup """
         if not self.wal_e or not self.wal_e_path:
             return False
@@ -187,7 +199,7 @@ class Postgresql:
         diff_in_bytes = long(backup_size)
         try:
             # get the difference in bytes between the current WAL location and the backup start offset
-            conn = psycopg2.connect(master_connurl)
+            conn = psycopg2.connect(master_connection)
             conn.autocommit = True
             cursor = conn.cursor()
             cursor.execute("SELECT pg_xlog_location_diff(pg_current_xlog_location(), %s)", (backup_start_lsn,))
@@ -216,10 +228,9 @@ class Postgresql:
             logger.error('Cannot start PostgreSQL because one is already running.')
             return False
 
-        pid_path = os.path.join(self.data_dir, 'postmaster.pid')
-        if os.path.exists(pid_path):
-            os.remove(pid_path)
-            logger.info('Removed %s', pid_path)
+        if os.path.exists(self.pid_path):
+            os.remove(self.pid_path)
+            logger.info('Removed %s', self.pid_path)
 
         ret = os.system(self._pg_ctl + ' start -o "{}"'.format(self.server_options())) == 0
         ret and self.load_replication_slots()
@@ -248,6 +259,9 @@ class Postgresql:
         return True
 
     def is_healthiest_node(self, cluster):
+        if self.is_leader():
+            return True
+
         if cluster.last_leader_operation - self.xlog_position() > self.config.get('maximum_lag_on_failover', 0):
             return False
 
@@ -255,18 +269,19 @@ class Postgresql:
             if member.hostname == self.name:
                 continue
             try:
-                member_conn = psycopg2.connect(member.address)
+                member_conn = psycopg2.connect(parseurl(member.address))
                 member_conn.autocommit = True
                 member_cursor = member_conn.cursor()
                 member_cursor.execute(
-                    "SELECT %s - (pg_last_xlog_replay_location() - '0/0000000'::pg_lsn)", (self.xlog_position(), ))
-                xlog_diff = member_cursor.fetchone()[0]
-                logger.info([self.name, member.hostname, xlog_diff])
+                    "SELECT pg_is_in_recovery(), %s - (pg_last_xlog_replay_location() - '0/0000000'::pg_lsn)",
+                    (self.xlog_position(), ))
+                row = member_cursor.fetchone()
                 member_cursor.close()
                 member_conn.close()
-                if xlog_diff < 0:
+                logger.error([self.name, member.hostname, row])
+                if not row[0] or row[1] < 0:
                     return False
-            except psycopg2.OperationalError:
+            except psycopg2.Error:
                 continue
         return True
 
@@ -283,7 +298,7 @@ class Postgresql:
     @staticmethod
     def primary_conninfo(leader_url):
         r = parseurl(leader_url)
-        return 'user={username} password={password} host={hostname} port={port} sslmode=prefer sslcompression=1'.format(**r)
+        return 'user={user} password={password} host={host} port={port} sslmode=prefer sslcompression=1'.format(**r)
 
     def check_recovery_conf(self, leader):
         if not os.path.isfile(self.recovery_conf):
@@ -357,7 +372,9 @@ primary_conninfo = '{}'
                 self.admin['username']), self.admin['password'])
 
     def xlog_position(self):
-        return self.query("SELECT pg_last_xlog_replay_location() - '0/0000000'::pg_lsn").fetchone()[0]
+        return self.query("""SELECT CASE WHEN pg_is_in_recovery()
+                                         THEN pg_last_xlog_replay_location() - '0/0000000'::pg_lsn
+                                         ELSE pg_current_xlog_location() - '0/00000'::pg_lsn END""").fetchone()[0]
 
     def load_replication_slots(self):
         cursor = self.query("SELECT slot_name FROM pg_replication_slots WHERE slot_type='physical'")
@@ -378,4 +395,4 @@ primary_conninfo = '{}'
         self.members = members
 
     def last_operation(self):
-        return self.query("SELECT pg_current_xlog_location() - '0/00000'::pg_lsn").fetchone()[0]
+        return self.xlog_position()

helpers/statuspage.py

@@ -1,65 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-import json
-import sys
-
-if sys.hexversion >= 0x03000000:
-    from http.server import BaseHTTPRequestHandler, HTTPServer
-else:
-    from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
-
-
-class StatusPage(BaseHTTPRequestHandler):
-
-    def do_GET(self):
-        content_type = 'text/plain'
-        if self.path == '/pg_master':
-            if not self.pg_is_in_recovery():
-                response, content = 200, 'I am currently a master'
-            else:
-                response, content = 503, 'I am not a master'
-        elif self.path == '/pg_slave':
-            if self.pg_is_in_recovery():
-                response, content = 200, 'I am currently a slave'
-            else:
-                response, content = 503, 'I am not a slave'
-        elif self.path == '/pg_status':
-            response, content = 200, self.pg_status()
-            content_type = 'application/json'
-        else:
-            response, content = 404, 'Page not found'
-
-        self.send_response(response)
-        self.send_header('Content-Type', content_type)
-        self.end_headers()
-        self.wfile.write(content.encode('utf-8'))
-
-    def pg_is_in_recovery(self):
-        cursor = self.server.postgresql.cursor()
-        cursor.execute('SELECT pg_is_in_recovery()')
-        res = cursor.fetchone()
-        return res[0]
-
-    def pg_status(self):
-        cursor = self.server.postgresql.cursor()
-        cursor.execute("""
-            SELECT pg_is_in_recovery(),
-                   to_char(pg_last_xact_replay_timestamp(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
-                   extract(epoch from now() - pg_last_xact_replay_timestamp()),
-                   inet_server_addr(),
-                   inet_server_port(),
-                   to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ')
-                    """)
-        res = cursor.fetchone()
-        status = {'role': ('master' if not res[0] else 'slave'), 'recovery': {'last_transaction_timestamp': res[1]},
-                  'server': {'hostaddr': res[3], 'port': res[4], 'start_time': res[5]}}
-
-        return json.dumps(status)
-
-
-def getHTTPServer(postgresql, http_port=8081, listen_address='0.0.0.0'):
-    server = HTTPServer((listen_address, http_port), StatusPage)
-    server.postgresql = postgresql
-
-    return server

postgres0.yml

@@ -1,5 +1,6 @@
 loop_wait: 10
-healthcheck_port: 8008
+restapi:
+  listen: 127.0.0.1:8008
 etcd:
   scope: batman
   ttl: 30

postgres1.yml

@@ -1,5 +1,6 @@
 loop_wait: 10
-healthcheck_port: 8009
+restapi:
+  listen: 127.0.0.1:8009
 etcd:
   scope: batman
   ttl: 30

tests/test_api.py

@@ -0,0 +1,66 @@
+import psycopg2
+import sys
+import unittest
+
+from helpers.api import RestApiHandler, RestApiServer
+from test_postgresql import psycopg2_connect
+
+if sys.hexversion >= 0x03000000:
+    from io import BytesIO as IO
+else:
+    from StringIO import StringIO as IO
+
+
+def false(*args, **kwargs):
+    return False
+
+
+def throws(*args, **kwargs):
+    raise psycopg2.OperationalError()
+
+
+class MockPostgresql:
+
+    def connection(self):
+        return psycopg2_connect()
+
+    def is_running(self):
+        return True
+
+
+class MockGovernor:
+
+    def __init__(self):
+        self.postgresql = MockPostgresql()
+
+
+class MockRequest:
+
+    def __init__(self, path):
+        self.path = path
+
+    def makefile(self, *args, **kwargs):
+        return IO(self.path)
+
+
+class MockRestApiServer(RestApiServer):
+
+    def __init__(self, Handler, path, *args):
+        self.governor = MockGovernor()
+        if len(args) > 0:
+            self.governor.postgresql.is_running = args[0]
+        self._cursor_holder = None
+        Handler(MockRequest(path), ('0.0.0.0', 8080), self)
+
+
+class TestRestApiHandler(unittest.TestCase):
+
+    def __init__(self, method_name='runTest'):
+        super(TestRestApiHandler, self).__init__(method_name)
+
+    def test_do_GET(self):
+        MockRestApiServer(RestApiHandler, b'GET /')
+        MockRestApiServer(RestApiHandler, b'GET /', throws)
+
+    def test_get_postgresql_status(self):
+        MockRestApiServer(RestApiHandler, b'GET /', false)

tests/test_etcd.py

@@ -17,28 +17,39 @@ class MockResponse:
         return json.loads(self.content)
 
 
+class MockPostgresql:
+    name = ''
+
+    def last_operation(self):
+        return 0
+
+
 def requests_get(url, **kwargs):
     if url.startswith('http://local'):
-        raise Exception()
+        raise requests.exceptions.RequestException()
     response = MockResponse()
     if url.startswith('http://remote'):
         response.content = '{"action":"get","node":{"key":"/service/batman5","dir":true,"nodes":[{"key":"/service/batman5/initialize","value":"postgresql0","modifiedIndex":1582,"createdIndex":1582},{"key":"/service/batman5/leader","value":"postgresql1","expiration":"2015-05-15T09:11:00.037397538Z","ttl":21,"modifiedIndex":20728,"createdIndex":20434},{"key":"/service/batman5/optime","dir":true,"nodes":[{"key":"/service/batman5/optime/leader","value":"2164261704","modifiedIndex":20729,"createdIndex":20729}],"modifiedIndex":20437,"createdIndex":20437},{"key":"/service/batman5/members","dir":true,"nodes":[{"key":"/service/batman5/members/postgresql1","value":"postgres://replicator:[email protected]:5434/postgres","expiration":"2015-05-15T09:10:59.949384522Z","ttl":21,"modifiedIndex":20727,"createdIndex":20727},{"key":"/service/batman5/members/postgresql0","value":"postgres://replicator:[email protected]:5433/postgres","expiration":"2015-05-15T09:11:09.611860899Z","ttl":30,"modifiedIndex":20730,"createdIndex":20730}],"modifiedIndex":1581,"createdIndex":1581}],"modifiedIndex":1581,"createdIndex":1581}}'
     elif url.startswith('http://other'):
         response.status_code = 404
+    elif url.startswith('http://noleader'):
+        response.content = '{"action":"get","node":{"key":"/service/batman5","dir":true,"nodes":[{"key":"/service/batman5/initialize","value":"postgresql0","modifiedIndex":1582,"createdIndex":1582},{"key":"/service/batman5/leader","value":"postgresql1","expiration":"2015-05-15T09:11:00.037397538Z","ttl":21,"modifiedIndex":20728,"createdIndex":20434},{"key":"/service/batman5/optime","dir":true,"nodes":[{"key":"/service/batman5/optime/leader","value":"2164261704","modifiedIndex":20729,"createdIndex":20729}],"modifiedIndex":20437,"createdIndex":20437},{"key":"/service/batman5/members","dir":true,"nodes":[{"key":"/service/batman5/members/postgresql0","value":"postgres://replicator:[email protected]:5433/postgres","expiration":"2015-05-15T09:11:09.611860899Z","ttl":30,"modifiedIndex":20730,"createdIndex":20730}],"modifiedIndex":1581,"createdIndex":1581}],"modifiedIndex":1581,"createdIndex":1581}}'
     return response
 
 
 def requests_put(url, **kwargs):
-    if url.startswith('http://local'):
-        raise Exception()
+    if url.startswith('http://local') or '/optime/leader' in url:
+        raise requests.exceptions.RequestException()
     response = MockResponse()
     response.status_code = 201
+    if url.startswith('http://other'):
+        response.status_code = 404
     return response
 
 
 def requests_delete(url):
     if url.startswith('http://local'):
-        raise Exception()
+        raise requests.exceptions.RequestException()
     response = MockResponse()
     response.status_code = 204
     return response
@@ -65,7 +76,7 @@ class TestEtcd(unittest.TestCase):
         self.assertRaises(Exception, self.etcd.get_client_path, '', 2)
 
     def test_put_client_path(self):
-        self.assertFalse(self.etcd.put_client_path(''))
+        self.assertRaises(EtcdError, self.etcd.put_client_path, '')
 
     def test_delete_client_path(self):
         self.assertFalse(self.etcd.delete_client_path(''))
@@ -77,6 +88,29 @@ class TestEtcd(unittest.TestCase):
         self.assertIsInstance(cluster, Cluster)
         self.etcd.base_client_url = self.etcd.base_client_url.replace('remote', 'other')
         self.etcd.get_cluster()
+        self.etcd.base_client_url = self.etcd.base_client_url.replace('other', 'noleader')
+        self.etcd.get_cluster()
 
     def test_current_leader(self):
         self.assertRaises(CurrentLeaderError, self.etcd.current_leader)
+
+    def test_touch_member(self):
+        self.assertFalse(self.etcd.touch_member('', ''))
+
+    def test_take_leader(self):
+        self.assertFalse(self.etcd.take_leader(''))
+
+    def test_attempt_to_acquire_leader(self):
+        self.assertFalse(self.etcd.attempt_to_acquire_leader(''))
+
+    def test_update_leader(self):
+        self.etcd.base_client_url = self.etcd.base_client_url.replace('local', 'remote')
+        self.assertTrue(self.etcd.update_leader(MockPostgresql()))
+        self.etcd.base_client_url = self.etcd.base_client_url.replace('remote', 'other')
+        self.assertFalse(self.etcd.update_leader(MockPostgresql()))
+
+    def test_race(self):
+        self.assertFalse(self.etcd.race('', ''))
+
+    def test_delete_member(self):
+        self.assertFalse(self.etcd.delete_member(''))

tests/test_governor.py

@@ -6,11 +6,16 @@ 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
 
+if sys.hexversion >= 0x03000000:
+    import http.server as BaseHTTPServer
+else:
+    import BaseHTTPServer
+
 
 def nop(*args, **kwargs):
     pass
@@ -20,6 +25,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,25 +37,37 @@ 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
         Postgresql.write_recovery_conf = nop
+        BaseHTTPServer.HTTPServer.__init__ = nop
 
     def tear_down(self):
         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:
@@ -61,7 +82,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)

tests/test_ha.py

@@ -72,7 +72,8 @@ class TestHa(unittest.TestCase):
         self.p = MockPostgresql()
         self.e = Etcd({'ttl': 30, 'host': 'remotehost', 'scope': 'test'})
         self.ha = Ha(self.p, self.e)
-        self.ha.cluster = Cluster(None, None, [])
+        self.ha.load_cluster_from_etcd()
+        self.ha.cluster = Cluster(False, None, None, [])
         self.ha.load_cluster_from_etcd = nop
 
     def test_start_as_slave(self):

tests/test_postgresql.py

@@ -22,37 +22,36 @@ def false(*args, **kwargs):
     return False
 
 
+def xlog_position():
+    return 1
+
+
 class MockCursor:
 
-    def __init__(self, server):
+    def __init__(self):
+        self.closed = False
         self.current = 0
         self.results = []
-        self.server = server
 
     def execute(self, sql, *params):
         if sql.startswith('blabla'):
             raise psycopg2.OperationalError()
+        elif sql.startswith('InterfaceError'):
+            raise psycopg2.InterfaceError()
         elif sql.startswith('SELECT slot_name'):
             self.results = [('blabla',), ('foobar',)]
         elif sql.startswith('SELECT pg_current_xlog_location()'):
-            self.results = [(0, )]
-        elif sql.startswith('SELECT %s - (pg_last_xlog_replay_location()'):
-            self.results = [(0, )]
-        elif sql.startswith('SELECT pg_last_xlog_replay_location()'):
-            self.results = [(0, )]
+            self.results = [(0,)]
+        elif sql.startswith('SELECT pg_is_in_recovery(), %s'):
+            if params[0][0] != 0:
+                raise psycopg2.OperationalError()
+            self.results = [(False, 0)]
+        elif sql.startswith('SELECT CASE WHEN pg_is_in_recovery()'):
+            self.results = [(0,)]
         elif sql.startswith('SELECT pg_is_in_recovery()'):
-            self.results = [(
-                self.server.mock_values['mock_recovery'],
-                None,
-                None,
-                None,
-                None,
-                None,
-                None,
-                None,
-                None,
-                None,
-            )]
+            self.results = [(False, )]
+        elif sql.startswith('SELECT to_char(pg_postmaster_start_time'):
+            self.results = [('', True, '', '', '', False)]
         else:
             self.results = [(
                 None,
@@ -82,14 +81,13 @@ class MockConnect:
 
     def __init__(self):
         self.autocommit = False
-        self.mock_values = {'mock_recovery': False}
+        self.closed = 0
 
     def cursor(self):
-        return MockCursor(self)
+        return MockCursor()
 
     def close(self):
-        if not self.autocommit:
-            raise psycopg2.OperationalError()
+        pass
 
 
 def psycopg2_connect(*args, **kwargs):
@@ -160,17 +158,25 @@ class TestPostgresql(unittest.TestCase):
 
     def test_query(self):
         self.p.query('select 1')
-        self.p.conn.autocommit = False
+        self.assertRaises(psycopg2.InterfaceError, self.p.query, 'InterfaceError')
+        self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
+        self.p._connection.closed = 2
+        self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
+        self.p._connection.closed = 2
+        self.p.disconnect = false
         self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
-        self.p.query('select %s', 1)
 
     def test_is_healthiest_node(self):
         leader = Member('leader', 'postgres://replicator:[email protected]:5435/postgres', 28)
         me = Member('test0', 'postgres://replicator:[email protected]:5434/postgres', 28)
         other = Member('test1', 'postgres://replicator:[email protected]:5433/postgres', 28)
-        cluster = Cluster(leader, 0, [leader, me, other])
+        cluster = Cluster(True, leader, 0, [me, other, leader])
         self.assertTrue(self.p.is_healthiest_node(cluster))
-        self.p.config['maximum_lag_on_failover'] = -1
+        self.p.is_leader = false
+        self.assertFalse(self.p.is_healthiest_node(cluster))
+        self.p.xlog_position = xlog_position
+        self.assertTrue(self.p.is_healthiest_node(cluster))
+        self.p.config['maximum_lag_on_failover'] = -2
         self.assertFalse(self.p.is_healthiest_node(cluster))
 
     def test_is_leader(self):

tests/test_statuspage.py

@@ -1,38 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-import unittest
-import sys
-
-from helpers.statuspage import StatusPage
-from test_postgresql import MockConnect
-
-if sys.hexversion >= 0x03000000:
-    from io import BytesIO as IO
-else:
-    from StringIO import StringIO as IO
-
-
-class TestStatusPage(unittest.TestCase):
-
-    def test_do_GET(self):
-        for mock_recovery in [True, False]:
-            for page in [b'GET /pg_master', b'GET /pg_slave', b'GET /pg_status', b'GET /not_found']:
-                self.http_server = MockServer(('0.0.0.0', 8888), StatusPage, page, mock_recovery)
-
-
-class MockRequest(object):
-
-    def __init__(self, path):
-        self.path = path
-
-    def makefile(self, *args, **kwargs):
-        return IO(self.path)
-
-
-class MockServer(object):
-
-    def __init__(self, ip_port, Handler, path, mock_recovery=False):
-        self.postgresql = MockConnect()
-        self.postgresql.mock_values['mock_recovery'] = mock_recovery
-        Handler(MockRequest(path), ip_port, self)