def test_disconnect_all_types(self):
"""Test that the disconnect method calls close on a all open connection types when no type is given"""
# Set up the test with mock data
connection_uri = 'someuri'
connection_type_1 = ConnectionType.DEFAULT
connection_type_2 = ConnectionType.EDIT
mock_connection_1 = MockConnection(dsn_parameters={
'host': 'myserver1',
'dbname': 'postgres1',
'user': '******'
})
mock_connection_2 = MockConnection(dsn_parameters={
'host': 'myserver2',
'dbname': 'postgres2',
'user': '******'
})
# Insert a ConnectionInfo object into the connection service's map
old_connection_details = ConnectionDetails.from_data({'abc': 123})
old_connection_info = ConnectionInfo(connection_uri,
old_connection_details)
old_connection_info.add_connection(connection_type_1,
mock_connection_1)
old_connection_info.add_connection(connection_type_2,
mock_connection_2)
self.connection_service.owner_to_connection_map[
connection_uri] = old_connection_info
# Close the connection by calling disconnect
response = self.connection_service._close_connections(
old_connection_info)
mock_connection_1.close.assert_called_once()
mock_connection_2.close.assert_called_once()
self.assertTrue(response)
def test_get_connection_info(self):
"""Test that get_connection_info returns the ConnectionInfo object corresponding to a connection"""
# Set up the test with mock data
connection_uri = 'someuri'
# Insert a ConnectionInfo object into the connection service's map
connection_details = ConnectionDetails.from_data({})
connection_info = ConnectionInfo(connection_uri, connection_details)
self.connection_service.owner_to_connection_map[connection_uri] = connection_info
# Get the connection info
actual_connection_info = self.connection_service.get_connection_info(connection_uri)
self.assertIs(actual_connection_info, connection_info)
def do_test():
# When I add context for 2 URIs with same connection details
operations_queue = OperationsQueue(self.mock_service_provider)
with mock.patch(COMPLETIONREFRESHER_PATH_PATH) as refresher_patch:
refresher_patch.return_value = self.refresher_mock
operations_queue.add_connection_context(self.connection_info)
conn_info2 = ConnectionInfo('newuri',
self.connection_info.details)
operations_queue.add_connection_context(conn_info2)
# Then I expect to only have connection
connect_mock: mock.MagicMock = self.mock_connection_service.connect
connect_mock.assert_called_once()
connect_params: ConnectRequestParams = connect_mock.call_args[
0][0]
self.assertEqual(connect_params.owner_uri,
self.expected_connection_uri)
def do_test():
# When I add context for 2 URIs with same connection details
operations_queue = OperationsQueue(self.mock_service_provider)
with mock.patch(COMPLETIONREFRESHER_PATH_PATH) as refresher_patch:
refresher_patch.return_value = self.refresher_mock
operations_queue.add_connection_context(self.connection_info)
conn_info2 = ConnectionInfo('newuri',
self.connection_info.details)
operations_queue.add_connection_context(conn_info2,
overwrite=True)
# Then I expect to only have 1 connection
# and I expect disconnect and reconnect to have been called
connect_mock: mock.MagicMock = self.mock_connection_service.connect
self.assertEqual(connect_mock.call_count, 2)
self.assertEqual(
connect_mock.call_args_list[0][0][0].owner_uri,
self.expected_connection_uri)
self.assertEqual(
connect_mock.call_args_list[1][0][0].owner_uri,
self.expected_connection_uri)
disconnect_mock: mock.MagicMock = self.mock_connection_service.disconnect
disconnect_mock.assert_called_once()
def test_get_connection_for_existing_connection(self):
"""Test that get_connection returns a connection that already exists for the given URI and type"""
# Set up the test with mock data
connection_uri = 'someuri'
connection_type = ConnectionType.EDIT
mock_connection = MockConnection(
dsn_parameters={
'host': 'myserver',
'dbname': 'postgres',
'user': '******'
})
# Insert a ConnectionInfo object into the connection service's map
connection_details = ConnectionDetails.from_data({})
connection_info = ConnectionInfo(connection_uri, connection_details)
self.connection_service.owner_to_connection_map[connection_uri] = connection_info
# Get the connection without first creating it
with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_connection)) as mock_psycopg2_connect:
connection = self.connection_service.get_connection(connection_uri, connection_type)
mock_psycopg2_connect.assert_called_once()
self.assertEqual(connection, mock_connection)
def setUp(self):
"""Set up the tests with a disaster recovery service and connection service with mock connection info"""
self.disaster_recovery_service = DisasterRecoveryService()
self.connection_service = ConnectionService()
self.task_service = TaskService()
self.disaster_recovery_service._service_provider = utils.get_mock_service_provider(
{
constants.CONNECTION_SERVICE_NAME: self.connection_service,
constants.TASK_SERVICE_NAME: self.task_service
})
# Create connection information for use in the tests
self.connection_details = ConnectionDetails()
self.host = 'test_host'
self.dbname = 'test_db'
self.username = '******'
self.connection_details.options = {
'host': self.host,
'dbname': self.dbname,
'user': self.username,
'port': 5432
}
self.test_uri = 'test_uri'
self.connection_info = ConnectionInfo(self.test_uri,
self.connection_details)
# Create backup parameters for the tests
self.request_context = utils.MockRequestContext()
self.backup_path = 'mock/path/test.sql'
self.backup_type = 'sql'
self.data_only = False
self.no_owner = True
self.schema = 'test_schema'
self.backup_params = BackupParams.from_dict({
'ownerUri': self.test_uri,
'backupInfo': {
'type': self.backup_type,
'path': self.backup_path,
'data_only': self.data_only,
'no_owner': self.no_owner,
'schema': self.schema
}
})
self.restore_path = 'mock/path/test.dump'
self.restore_params = RestoreParams.from_dict({
'ownerUri': self.test_uri,
'options': {
'path': self.restore_path,
'data_only': self.data_only,
'no_owner': self.no_owner,
'schema': self.schema
}
})
self.pg_dump_exe = 'pg_dump'
self.pg_restore_exe = 'pg_restore'
# Create the mock task for the tests
self.mock_action = mock.Mock()
self.mock_task = Task(None, None, None, None, None,
self.request_context, self.mock_action)
self.mock_task.start = mock.Mock()
