improve test coverage

patroni/postgresql.py

@@ -509,8 +509,7 @@ recovery_target_timeline = 'latest'
                 result = {l.split(':')[0].replace('Current ', '', 1): l.split(':')[1].strip() for l in data if l}
         except subprocess.CalledProcessError:
             logger.exception("Error when calling pg_controldata")
-        finally:
-            return result
+        return result
 
     def read_postmaster_opts(self):
         """ returns the list of option names/values from postgres.opts, Empty dict if read failed or no file """

patroni/scripts/aws.py

@@ -66,6 +66,7 @@ class AWSConnection(object):
 
 
 def main():
+    logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO)
     if len(sys.argv) == 4 and sys.argv[1] in ('on_start', 'on_stop', 'on_role_change'):
         AWSConnection(cluster_name=sys.argv[3]).on_role_change(sys.argv[2])
     else:

patroni/scripts/wale_restore.py

@@ -1,4 +1,4 @@
-#!/usr/bin/python
+#!/usr/bin/env python
 
 # sample script to clone new replicas using WAL-E restore
 # falls back to pg_basebackup if WAL-E restore fails, or if
@@ -36,7 +36,6 @@ import argparse
 if sys.hexversion >= 0x03000000:
     long = int
 
-logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
@@ -140,6 +139,7 @@ class WALERestore(object):
 
 
 def main():
+    logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO)
     parser = argparse.ArgumentParser(description='Script to image replicas using WAL-E')
     parser.add_argument('--scope', required=True)
     parser.add_argument('--role', required=False)

tests/test_aws.py

@@ -1,9 +1,11 @@
-import unittest
-import requests
 import boto.ec2
+import requests
+import sys
+import unittest
 
+from mock import Mock, patch
 from collections import namedtuple
-from patroni.scripts.aws import AWSConnection
+from patroni.scripts.aws import AWSConnection, main as _main
 from requests.exceptions import RequestException
 
 
@@ -84,3 +86,9 @@ class TestAWSConnection(unittest.TestCase):
     def test_aws_tag_ec2_error(self):
         self.error = True
         self.assertFalse(self.conn._tag_ec2("master"))
+
+    @patch('sys.exit', Mock())
+    def test_main(self):
+        self.assertIsNone(_main())
+        sys.argv = ['aws.py', 'on_start', 'replica', 'foo']
+        self.assertIsNone(_main())

tests/test_patroni.py

@@ -7,7 +7,7 @@ from mock import Mock, patch
 from patroni.api import RestApiServer
 from patroni.async_executor import AsyncExecutor
 from patroni.etcd import Etcd
-from patroni import Patroni, main
+from patroni import Patroni, main as _main
 from patroni.zookeeper import ZooKeeper
 from six.moves import BaseHTTPServer
 from test_etcd import Client, SleepException, etcd_read, etcd_write
@@ -51,14 +51,14 @@ class TestPatroni(unittest.TestCase):
     @patch.object(Etcd, 'delete_leader', Mock())
     @patch.object(Client, 'machines')
     def test_patroni_main(self, mock_machines):
-        main()
+        _main()
         sys.argv = ['patroni.py', 'postgres0.yml']
 
         mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
         with patch.object(Patroni, 'run', Mock(side_effect=SleepException())):
-            self.assertRaises(SleepException, main)
+            self.assertRaises(SleepException, _main)
         with patch.object(Patroni, 'run', Mock(side_effect=KeyboardInterrupt())):
-            main()
+            _main()
 
     @patch('time.sleep', Mock(side_effect=SleepException()))
     def test_run(self):

tests/test_postgresql.py

@@ -254,6 +254,7 @@ class TestPostgresql(unittest.TestCase):
     @patch('patroni.postgresql.Postgresql.remove_data_directory', MagicMock(return_value=True))
     @patch('patroni.postgresql.Postgresql.single_user_mode', MagicMock(return_value=1))
     @patch('patroni.postgresql.Postgresql.write_pgpass', MagicMock(return_value=dict()))
+    @patch('subprocess.check_output', Mock(return_value=0, side_effect=pg_controldata_string))
     def test_follow(self, mock_pg_rewind):
         self.p.follow(None)
         self.p.follow(self.leader)
@@ -274,6 +275,7 @@ class TestPostgresql(unittest.TestCase):
         with mock.patch('patroni.postgresql.Postgresql.check_recovery_conf', MagicMock(return_value=True)):
             self.assertTrue(self.p.follow(None))
 
+    @patch('subprocess.check_output', Mock(return_value=0, side_effect=pg_controldata_string))
     def test_can_rewind(self):
         tmp = self.p.pg_rewind
         self.p.pg_rewind = None
@@ -283,7 +285,7 @@ class TestPostgresql(unittest.TestCase):
             self.assertFalse(self.p.can_rewind)
         with mock.patch('subprocess.call', side_effect=OSError("foo")):
             self.assertFalse(self.p.can_rewind)
-        tmp = self.p.controldata()
+        tmp = self.p.controldata
         self.p.controldata = lambda: {'wal_log_hints setting': 'on'}
         self.assertTrue(self.p.can_rewind)
         self.p.controldata = tmp
@@ -390,22 +392,16 @@ class TestPostgresql(unittest.TestCase):
             self.p.remove_data_directory()
         self.p.remove_data_directory()
 
-    @patch('subprocess.check_output', MagicMock(return_value=0, side_effect=pg_controldata_string))
-    @patch('subprocess.check_output', side_effect=subprocess.CalledProcessError)
-    @patch('subprocess.check_output', side_effect=Exception('Failed'))
-    def test_controldata(self, check_output_call_error, check_output_generic_exception):
-        data = self.p.controldata()
-        self.assertEquals(len(data), 50)
-        self.assertEquals(data['Database cluster state'], 'shut down in recovery')
-        self.assertEquals(data['wal_log_hints setting'], 'on')
-        self.assertEquals(int(data['Database block size']), 8192)
+    def test_controldata(self):
+        with patch('subprocess.check_output', Mock(return_value=0, side_effect=pg_controldata_string)):
+            data = self.p.controldata()
+            self.assertEquals(len(data), 50)
+            self.assertEquals(data['Database cluster state'], 'shut down in recovery')
+            self.assertEquals(data['wal_log_hints setting'], 'on')
+            self.assertEquals(int(data['Database block size']), 8192)
 
-        subprocess.check_output = check_output_call_error
-        data = self.p.controldata()
-        self.assertEquals(data, dict())
-
-        subprocess.check_output = check_output_generic_exception
-        self.assertRaises(Exception, self.p.controldata())
+        with patch('subprocess.check_output', Mock(side_effect=subprocess.CalledProcessError(1, ''))):
+            self.assertEquals(self.p.controldata(), {})
 
     def test_read_postmaster_opts(self):
         m = mock_open(read_data=postmaster_opts_string())

tests/test_wale_restore.py

@@ -1,8 +1,10 @@
-import unittest
-from mock import MagicMock, patch, PropertyMock
 import psycopg2
 import subprocess
-from patroni.scripts.wale_restore import WALERestore, main
+import sys
+import unittest
+
+from mock import MagicMock, patch, PropertyMock
+from patroni.scripts.wale_restore import WALERestore, main as _main
 
 
 def fake_cursor_fetchone(*args, **kwargs):
@@ -94,6 +96,8 @@ class TestWALERestore(unittest.TestCase):
             with patch.object(self.wale_restore, 'create_replica_with_s3', MagicMock(return_value=0)):
                 self.assertEqual(self.wale_restore.run(), 0)
 
+    @patch('sys.exit', MagicMock())
+    @patch.object(WALERestore, 'run', MagicMock(return_value=0))
     def test_main(self):
-        with patch('sys.exit', MagicMock(return_value=0)):
-            self.assertEqual(main(), None)
+        self.assertEqual(_main(), None)
+        sys.argv