diff --git a/tests/units/fastsync/commons/test_fastsync_tap_postgres.py b/tests/units/fastsync/commons/test_fastsync_tap_postgres.py index 7dafa9b34..dee154041 100644 --- a/tests/units/fastsync/commons/test_fastsync_tap_postgres.py +++ b/tests/units/fastsync/commons/test_fastsync_tap_postgres.py @@ -21,11 +21,6 @@ def setUp(self) -> None: self.postgres.executed_queries_primary_host = [] self.postgres.executed_queries = [] - def primary_host_query_mock(query, _=None): - self.postgres.executed_queries_primary_host.append(query) - - self.postgres.primary_host_query = primary_host_query_mock - def test_generate_repl_slot_name(self): """Validate if the replication slot name generated correctly""" # Provide only database name @@ -70,17 +65,20 @@ def test_create_replication_slot_1(self): Validate if replication slot creation SQL commands generated correctly in case no v15 slots exists """ - def execute_mock(query): + def execute_mock(query, _=None): print('Mocked execute called') self.postgres.executed_queries_primary_host.append(query) # mock cursor with execute method cursor_mock = MagicMock().return_value cursor_mock.__enter__.return_value.execute.side_effect = execute_mock - type(cursor_mock.__enter__.return_value).rowcount = PropertyMock(return_value=0) + # First query -> 0 rows, second query -> 1 row + type(cursor_mock.__enter__.return_value).rowcount = PropertyMock(side_effect=[0, 1]) + cursor_mock.__enter__.return_value.fetchall.return_value = [{'lsn': '24368/44107178'}] # mock PG connection instance with ability to open cursor - pg_con = Mock() + pg_con = MagicMock() + pg_con.__enter__.return_value = pg_con pg_con.cursor.return_value = cursor_mock self.postgres.primary_host_conn = pg_con @@ -96,17 +94,20 @@ def test_create_replication_slot_2(self): Validate if replication slot creation SQL commands generated correctly in case a v15 slots exists """ - def execute_mock(query): + def execute_mock(query, _=None): print('Mocked execute called') self.postgres.executed_queries_primary_host.append(query) # mock cursor with execute method cursor_mock = MagicMock().return_value cursor_mock.__enter__.return_value.execute.side_effect = execute_mock - type(cursor_mock.__enter__.return_value).rowcount = PropertyMock(return_value=1) + # First query -> 1 row, second query -> 1 row + type(cursor_mock.__enter__.return_value).rowcount = PropertyMock(side_effect=[1, 1]) + cursor_mock.__enter__.return_value.fetchall.return_value = [{'lsn': '24368/44107178'}] # mock PG connection instance with ability to open cursor - pg_con = Mock() + pg_con = MagicMock() + pg_con.__enter__.return_value = pg_con pg_con.cursor.return_value = cursor_mock self.postgres.primary_host_conn = pg_con