diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 7fab3d65..4e6c0f0e 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -14,7 +14,7 @@ from patroni.postgresql import Postgresql, STATE_REJECT, STATE_NO_RESPONSE from patroni.postmaster import PostmasterProcess from patroni.utils import RetryFailedError from six.moves import builtins -from threading import Thread +from threading import Thread, current_thread class MockCursor(object): @@ -790,10 +790,12 @@ class TestPostgresql(unittest.TestCase): def test_wait_for_startup(self): state = {'sleeps': 0, 'num_rejects': 0, 'final_return': 0} + self.__thread_ident = current_thread().ident def increment_sleeps(*args): - print("Sleep") - state['sleeps'] += 1 + if current_thread().ident == self.__thread_ident: + print("Sleep") + state['sleeps'] += 1 def isready_return(*args): ret = 1 if state['sleeps'] < state['num_rejects'] else state['final_return'] diff --git a/tests/test_wale_restore.py b/tests/test_wale_restore.py index 57fb7509..057cfd96 100644 --- a/tests/test_wale_restore.py +++ b/tests/test_wale_restore.py @@ -7,6 +7,7 @@ from patroni.scripts import wale_restore from patroni.scripts.wale_restore import WALERestore, main as _main, get_major_version from six.moves import builtins from test_postgresql import MockConnect, psycopg2_connect +from threading import current_thread wale_output_header = ( @@ -42,6 +43,13 @@ class TestWALERestore(unittest.TestCase): '/etc', 100, 100, 1, 0, WALE_TEST_RETRIES) def test_should_use_s3_to_create_replica(self): + self.__thread_ident = current_thread().ident + sleeps = [0] + + def mock_sleep(*args): + if current_thread().ident == self.__thread_ident: + sleeps[0] += 1 + self.assertTrue(self.wale_restore.should_use_s3_to_create_replica()) with patch.object(MockConnect, 'server_version', PropertyMock(return_value=100000)): self.assertTrue(self.wale_restore.should_use_s3_to_create_replica()) @@ -55,13 +63,11 @@ class TestWALERestore(unittest.TestCase): self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) - with patch('time.sleep', Mock(return_value=None)) as mock_sleep: + with patch('time.sleep', mock_sleep): self.wale_restore.no_master = 1 - assert self.wale_restore.should_use_s3_to_create_replica() + self.assertTrue(self.wale_restore.should_use_s3_to_create_replica()) # verify retries - mock_sleep.assert_has_calls( - [((wale_restore.RETRY_SLEEP_INTERVAL,),)] * WALE_TEST_RETRIES - ) + self.assertEqual(sleeps[0], WALE_TEST_RETRIES) self.wale_restore.master_connection = '' self.assertTrue(self.wale_restore.should_use_s3_to_create_replica()) @@ -104,13 +110,20 @@ class TestWALERestore(unittest.TestCase): @patch('sys.exit', Mock()) def test_main(self): + self.__thread_ident = current_thread().ident + sleeps = [0] + + def mock_sleep(*args): + if current_thread().ident == self.__thread_ident: + sleeps[0] += 1 + with patch.object(WALERestore, 'run', Mock(return_value=0)): self.assertEqual(_main(), 0) with patch.object(WALERestore, 'run', Mock(return_value=1)), \ - patch('time.sleep', Mock(return_value=None)) as mock_sleep: + patch('time.sleep', mock_sleep): self.assertEqual(_main(), 1) - assert mock_sleep.call_count == WALE_TEST_RETRIES + self.assertTrue(sleeps[0], WALE_TEST_RETRIES) @patch('os.path.isfile', Mock(return_value=True)) def test_get_major_version(self):