Add mock-based unit tests for the restore module.

- add unit tests for the restore module
- additional dependencies in requirement
- harden the code that calls restore callbacks
- remove WAL-E related code from postgresql.py

helpers/postgresql.py

@@ -61,10 +61,6 @@ class Postgresql:
         self.is_promoted = False
 
         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.local_address = self.get_local_address()
         connect_address = config.get('connect_address', None) or self.local_address
@@ -152,9 +148,13 @@ class Postgresql:
     def create_replica(self, master_connection, env):
         connstring = self.build_connstring(master_connection, master_connection)
         cmd = self.config['restore']
-        ret = subprocess.call(shlex.split(os.path.abspath(cmd))+[self.scope, "replica",
-                              self.data_dir, connstring], env=env)
-        self.delete_trigger_file()
+        try:
+            ret = subprocess.call(shlex.split(os.path.abspath(cmd))+[self.scope, "replica",
+                                  self.data_dir, connstring], env=env)
+            self.delete_trigger_file()
+        except Exception as e:
+            logger.error("Error when creating replica: {0}".format(e))
+            return 1
         return ret
 
     def is_leader(self, check_only=False):

requirements-py2.txt

@@ -1,6 +1,8 @@
 boto
 dnspython
+mock
 psycopg2
 PyYAML
 requests
+six >= 1.7
 kazoo>=2.2.1

requirements-py3.txt

@@ -1,4 +1,5 @@
 boto
+mock
 dnspython3
 psycopg2
 PyYAML

scripts/restore.py

@@ -17,17 +17,21 @@ import psycopg2
 import subprocess
 import sys
 
+
+if sys.hexversion >= 0x03000000:
+    long = int
+
 logger = logging.getLogger(__name__)
 
 
 class Restore(object):
 
-    def __init__(self, scope, role, datadir, connstring):
+    def __init__(self, scope, role, datadir, connstring, env=None):
         self.scope = scope
         self.role = role
         self.master_connection = Restore.parse_connstring(connstring)
         self.data_dir = datadir
-        self.env = os.environ.copy()
+        self.env = os.environ.copy() if not env else env
 
     @staticmethod
     def parse_connstring(connstring):
@@ -41,6 +45,9 @@ class Restore(object):
                 result[key.strip()] = val.strip()
         return result
 
+    def setup(self):
+        pass
+
     def replica_method(self):
         return self.create_replica_with_pg_basebackup
 
@@ -50,26 +57,31 @@ class Restore(object):
     def run(self):
         """ creates a new replica using either pg_basebackup or WAL-E """
         method_fn = self.replica_method()
-        ret = method_fn()
+        ret = method_fn() if method_fn else 1
         if ret != 0 and self.replica_fallback_method() is not None:
             ret = (self.replica_fallback_method())()
         return ret
 
     def create_replica_with_pg_basebackup(self):
-        ret = subprocess.call(['pg_basebackup', '-R', '-D',
-                               self.data_dir, '--host=' + self.master_connection['host'],
-                               '--port=' + str(self.master_connection['port']),
-                               '-U', self.master_connection['user']],
-                              env=self.env)
+        try:
+            ret = subprocess.call(['pg_basebackup', '-R', '-D',
+                                   self.data_dir, '--host=' + self.master_connection['host'],
+                                   '--port=' + str(self.master_connection['port']),
+                                   '-U', self.master_connection['user']],
+                                  env=self.env)
+        except Exception as e:
+            logger.error('Error when fetching backup with pg_basebackup: {0}'.format(e))
+            return 1
         return ret
 
-
 class WALERestore(Restore):
 
-    def __init__(self, scope, role, datadir, connstring):
-        super(WALERestore, self).__init__(scope, role, datadir, connstring)
+    def __init__(self, scope, role, datadir, connstring, env=None):
+        super(WALERestore, self).__init__(scope, role, datadir, connstring, env)
         # check the environment variables
         self.init_error = False
+
+    def setup(self):
         if (self.env.get('WAL_S3_BUCKET') and
             self.env.get('WALE_BACKUP_THRESHOLD_PERCENTAGE') and
            self.env.get('WALE_BACKUP_THRESHOLD_MEGABYTES')) is None:
@@ -105,9 +117,9 @@ class WALERestore(Restore):
                 self.init_error = True
 
     def replica_method(self):
-        if self.should_use_s3_to_create_replica(self):
+        if self.should_use_s3_to_create_replica():
             return self.create_replica_with_s3
-        return 1
+        return None
 
     def replica_fallback_method(self):
         return self.create_replica_with_pg_basebackup
@@ -185,9 +197,11 @@ class WALERestore(Restore):
     def create_replica_with_s3(self):
         if self.init_error:
             return 1
-
-        ret = subprocess.call(self.wal_e.cmd + ' backup-fetch {} LATEST'.format(self.data_dir), env=self.env)
-        self.restore_configuration_files()
+        try:
+            ret = subprocess.call(self.wal_e.cmd + ' backup-fetch {} LATEST'.format(self.data_dir), env=self.env)
+        except Exception as e:
+            logger.error('Error when fetching backup with WAL-E: {0}'.format(e))
+            return 1
         return ret
 
 
@@ -195,5 +209,6 @@ if __name__ == '__main__':
     if len(sys.argv) == 5:
         # scope, role, datadir, connstring
         restore = WALERestore(*(sys.argv[1:]))
+        restore.setup()
         sys.exit(restore.run())
     sys.exit("Usage: {0} scope role datadir connstring".format(sys.argv[0]))

tests/test_restore.py

@@ -0,0 +1,114 @@
+import unittest
+from mock import MagicMock, patch
+import os
+from scripts.restore import Restore, WALERestore
+
+
+def fake_cursor_fetchone(*args, **kwargs):
+    return ('16777216',)
+
+
+def fake_call_fail_for_wal_e(*args, **kwargs):
+    if len(args) > 0 and 'backup-fetch' in args[0]:
+        return 1
+    return 0
+
+
+def fake_call_fail_for_base_backup(*args, **kwargs):
+    if len(args) > 0 and 'backup-fetch' in args[0]:
+        return 0
+    return 1
+
+
+def fake_backup_data(self, *args, **kwargs):
+    """ return the fake result of WAL-E backup-list"""
+    return """name    last_modified   expanded_size_bytes wal_segment_backup_start    wal_segment_offset_backup_start wal_segment_backup_stop wal_segment_offset_backup_stop
+base_00000001000000000000007F_00000040  2015-05-18T10:13:25.000Z 167772160   00000001000000000000007F    00000040 00000001000000000000007F    00000240
+"""
+
+
+class TestRestore(unittest.TestCase):
+
+    def setUp(self):
+        self.restore = Restore("batman", "master", "/data", "host=batman port=5432 user=batman")
+        pass
+
+    def tearDown(self):
+        pass
+
+    def test_parse_connstring(self):
+        self.assertDictEqual(self.restore.master_connection, {'host': 'batman', 'port': '5432', 'user': 'batman'})
+
+    @patch('subprocess.call', MagicMock(return_value=0))
+    def test_run(self):
+        ret = self.restore.run()
+        self.assertEqual(ret, 0)
+
+    @patch('subprocess.call', MagicMock(return_value=1))
+    def test_run_fail(self):
+        ret = self.restore.run()
+        self.assertEqual(ret, 1)
+
+
+
+@patch('os.access', MagicMock(return_value=True))
+@patch('os.makedirs', MagicMock(return_value=True))
+@patch('os.path.exists', MagicMock(return_value=True))
+@patch('os.path.isdir', MagicMock(return_value=True))
+@patch('psycopg2.extensions.cursor.fetchone', MagicMock(side_effect=fake_cursor_fetchone))
+@patch('psycopg2.extensions.cursor', MagicMock(autospec=True))
+@patch('psycopg2.extensions.connection', MagicMock(autospec=True))
+@patch('psycopg2.connect', MagicMock(autospec=True))
+@patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data))
+class TestWALERestore(unittest.TestCase):
+
+    def setUp(self):
+        env = {}
+        env['WAL_S3_BUCKET'] = 'batman'
+        env['WALE_BACKUP_THRESHOLD_PERCENTAGE'] = 100
+        env['WALE_BACKUP_THRESHOLD_MEGABYTES'] = 100
+        self.wale_restore = WALERestore("batman", "master", "/data", "host=batman port=5432 user=batman", env=env)
+
+    def tearDown(self):
+        pass
+
+    def test_setup(self):
+        self.wale_restore.setup()
+        self.assertFalse(self.wale_restore.init_error)
+
+
+    # have to redefine the class-level os.access mock inside the function
+    # since the class-level mock will be applied after the function level one.
+    @patch('os.access', return_value=False)
+    def test_setup_fail(self, mock_no_access):
+        os.access = mock_no_access
+        self.wale_restore.setup()
+        self.assertTrue(self.wale_restore.init_error)
+
+
+    # The 3 tests above only differ with the mock function instead of a subprocess call
+    # in the first one, subprocess call should return success only for wal-e command,
+    # checking the primary use-case of restoring from WAL-E backup.
+    # In the second one, we test fallbacks by failing at WAL-E, but succeeding at
+    # pg_basebackup.
+    # Finally, the last use case is when all subprocess.call fails. resulting in a
+    # failure to restore from replica
+    @patch('subprocess.call',
+           MagicMock(side_effect=lambda *args, **kwargs: 0 if 'wal-e' in args[0] else 1))
+    def test_run(self):
+        self.wale_restore.setup()
+        ret = self.wale_restore.run()
+        self.assertEqual(ret, 0)
+
+    @patch('subprocess.call',
+           MagicMock(side_effect=lambda *args, **kwargs: 0 if 'pg_basebackup' in args[0] else 1))
+    def test_run_fallback(self):
+        self.wale_restore.setup()
+        ret = self.wale_restore.run()
+        self.assertEqual(ret, 0)
+
+    @patch('subprocess.call', MagicMock(return_value=1))
+    def test_run_all_fail(self):
+        self.wale_restore.setup()
+        ret = self.wale_restore.run()
+        self.assertEqual(ret, 1)