def setUp(self):
"""Set up the tests with common connection parameters"""
# Set up the mock connection service and connection info
self.connection_service = ConnectionService()
self.connection_service._service_provider = utils.get_mock_service_provider(
{WORKSPACE_SERVICE_NAME: WorkspaceService()})
self.owner_uri = 'test_uri'
self.connection_type = ConnectionType.DEFAULT
self.connect_params: ConnectRequestParams = ConnectRequestParams.from_dict(
{
'ownerUri': self.owner_uri,
'type': self.connection_type,
'connection': {
'options': {}
}
})
self.mock_psycopg_connection = MockPsycopgConnection(dsn_parameters={
'host': 'myserver',
'dbname': 'postgres',
'user': '******'
})
# Mock psycopg2's connect method to store the current cancellation token. This lets us
# capture the cancellation token state as it would be during a long-running connection.
self.token_store = []
def test_connect(self):
"""Test that the service connects to a MySQL server"""
# Set up the parameters for the connection
params: ConnectRequestParams = ConnectRequestParams.from_dict({
'ownerUri':
'someUri',
'type':
ConnectionType.DEFAULT,
'connection': {
'options': {
'user': '******',
'password': '******',
'host': 'myserver',
'dbname': 'mysql'
}
}
})
# Set up the connection service and call its connect method with the supported options
with mock.patch(
'pymysql.connect',
new=mock.Mock(return_value=self.mock_pymysql_connection)):
response = self.connection_service.connect(params)
# Verify that pymysql's connection method was called and that the
# response has a connection id, indicating success.
self.assertIs(
self.connection_service.owner_to_connection_map[
params.owner_uri].get_connection(params.type)._conn,
self.mock_pymysql_connection)
self.assertIsNotNone(response.connection_id)
self.assertIsNotNone(response.server_info.server_version)
self.assertFalse(response.server_info.is_cloud)
def test_non_default_database(self):
"""Test that if a database is given, the default database is not used"""
# Set up the connection params and default database name
default_db = 'test_db'
actual_db = 'postgres'
self.connection_service._service_provider[
WORKSPACE_SERVICE_NAME].configuration.pgsql.default_database = default_db
params: ConnectRequestParams = ConnectRequestParams.from_dict({
'ownerUri':
'someUri',
'type':
ConnectionType.DEFAULT,
'connection': {
'options': {
'user': '******',
'password': '******',
'host': 'myserver',
'dbname': actual_db
}
}
})
# If I connect with an empty database name
with mock.patch('ossdbtoolsservice.connection.connection_service._build_connection_response'), \
mock.patch('psycopg2.connect', return_value=MockPsycopgConnection()) as mock_psycopg2_connect:
self.connection_service.connect(params)
# Then psycopg2's connect method was called with the default database
calls = mock_psycopg2_connect.mock_calls
self.assertEqual(len(calls), 1)
self.assertNotEqual(calls[0][2]['dbname'], default_db)
self.assertEqual(calls[0][2]['dbname'], actual_db)
def test_changing_options_disconnects_existing_connection(self):
"""
Test that the connect method disconnects an existing connection when trying to open the same connection with
different options
"""
# Set up the test with mock data
connection_uri = 'someuri'
connection_type = ConnectionType.DEFAULT
mock_connection = MockPGServerConnection(cur=None,
host='myserver',
name='postgres',
user='******')
# Insert a ConnectionInfo object into the connection service's map
old_connection_details = ConnectionDetails.from_data({
'host': 'myserver',
'dbname': 'postgres',
'user': '******',
'abc': 123
})
old_connection_info = ConnectionInfo(connection_uri,
old_connection_details)
old_connection_info.add_connection(connection_type, mock_connection)
self.connection_service.owner_to_connection_map[
connection_uri] = old_connection_info
# Create a different request with the same owner uri
params: ConnectRequestParams = ConnectRequestParams.from_dict({
'ownerUri':
connection_uri,
'type':
connection_type,
'connection': {
'options': {
'host': 'myserver',
'dbname': 'postgres',
'user': '******',
'abc': 234
}
}
})
# Connect with different options, and verify that disconnect was called
with mock.patch('psycopg2.connect',
new=mock.Mock(return_value=mock_connection)):
self.connection_service.connect(params)
mock_connection.close.assert_called_once()
def test_same_options_uses_existing_connection(self):
"""Test that the connect method uses an existing connection when connecting again with the same options"""
# Set up the test with mock data
connection_uri = 'someuri'
connection_type = ConnectionType.DEFAULT
mock_psycopg_connection = MockPsycopgConnection(dsn_parameters={
'host': 'myserver',
'dbname': 'postgres',
'user': '******'
})
mock_server_connection = MockPGServerConnection(cur=None,
host='myserver',
name='postgres',
user='******')
# Insert a ConnectionInfo object into the connection service's map
old_connection_details = ConnectionDetails.from_data({
'host': 'myserver',
'dbname': 'postgres',
'user': '******',
'abc': 123
})
old_connection_info = ConnectionInfo(connection_uri,
old_connection_details)
old_connection_info.add_connection(connection_type,
mock_server_connection)
self.connection_service.owner_to_connection_map[
connection_uri] = old_connection_info
# Connect with identical options, and verify that disconnect was not called
params: ConnectRequestParams = ConnectRequestParams.from_dict({
'ownerUri':
connection_uri,
'type':
connection_type,
'connection': {
'options': old_connection_details.options
}
})
with mock.patch('psycopg2.connect',
new=mock.Mock(return_value=mock_psycopg_connection)
) as mock_psycopg2_connect:
response = self.connection_service.connect(params)
mock_psycopg2_connect.assert_not_called()
mock_psycopg_connection.close.assert_not_called()
self.assertIsNotNone(response.connection_id)
def test_connect_with_access_token(self):
"""Test that the service connects to a MySQL server using an access token as a password"""
# Set up the parameters for the connection
params: ConnectRequestParams = ConnectRequestParams.from_dict({
'ownerUri':
'someUri',
'type':
ConnectionType.DEFAULT,
'connection': {
'options': {
'user': '******',
'azureAccountToken': 'exampleToken',
'host': 'myserver',
'dbname': 'mysql'
}
}
})
# Set up pymysql instance for connection service to call
mock_connect_method = mock.Mock(
return_value=self.mock_pymysql_connection)
# Set up the connection service and call its connect method with the supported options
with mock.patch('pymysql.connect', new=mock_connect_method):
response = self.connection_service.connect(params)
# Verify that pymysql's connection method was called with password set to account token.
mock_connect_method.assert_called_once_with(
user='******',
password='******',
host='myserver',
port=DEFAULT_PORT[MYSQL_PROVIDER_NAME],
database='mysql')
# Verify that pymysql's connection method was called and that the
# response has a connection id, indicating success.
self.assertIs(
self.connection_service.owner_to_connection_map[
params.owner_uri].get_connection(params.type)._conn,
self.mock_pymysql_connection)
self.assertIsNotNone(response.connection_id)
self.assertIsNotNone(response.server_info.server_version)
self.assertFalse(response.server_info.is_cloud)
def test_handle_connect_request(self):
"""Test that the handle_connect_request method kicks off a new thread to do the connection"""
# Setup: Create a mock request context to handle output
rc = utils.MockRequestContext()
connect_response = ConnectionCompleteParams()
self.connection_service.connect = Mock(return_value=connect_response)
# If: I make a request to connect
params: ConnectRequestParams = ConnectRequestParams.from_dict({
'ownerUri':
'someUri',
'type':
ConnectionType.QUERY,
'connection': {
'server_name': 'someserver',
'user_name': 'someuser',
'database_name': 'somedb',
'options': {
'password': '******'
}
}
})
# Connect and wait for the thread to finish executing, then verify the connection information
self.connection_service.handle_connect_request(rc, params)
connection_thread = self.connection_service.owner_to_thread_map[
params.owner_uri]
self.assertIsNotNone(connection_thread)
connection_thread.join()
# Then:
# ... Connect should have been called once
self.connection_service.connect.assert_called_once_with(params)
# ... A True should have been sent as the response to the request
rc.send_response.assert_called_once_with(True)
# ... A connection complete notification should have been sent back as well
rc.send_notification.assert_called_once_with(
CONNECTION_COMPLETE_METHOD, connect_response)
# ... An error should not have been called
rc.send_error.assert_not_called()
def test_response_when_connect_fails(self):
"""Test that the proper response is given when a connection fails"""
error_message = 'some error'
params: ConnectRequestParams = ConnectRequestParams.from_dict({
'ownerUri':
'someUri',
'type':
ConnectionType.DEFAULT,
'connection': {
'options': {
'connectionString': ''
}
}
})
with mock.patch('psycopg2.connect',
new=mock.Mock(side_effect=Exception(error_message))):
response = self.connection_service.connect(params)
# The response should not have a connection ID and should contain the error message
self.assertIsNone(response.connection_id)
self.assertEqual(response.error_message, error_message)
