def test_generate_uri_missing_params(self):
# Setup: Create the parameter sets that will be missing a param each
params = [
ConnectionDetails.from_data({
'host': None,
'dbname': TEST_DBNAME,
'user': TEST_USER
}),
ConnectionDetails.from_data({
'host': TEST_HOST,
'dbname': None,
'user': TEST_USER
}),
ConnectionDetails.from_data({
'host': TEST_HOST,
'dbname': TEST_DBNAME,
'user': None
})
]
for param_set in params:
# If: I generate a session URI from params that are missing a value
# Then: I should get an exception
with self.assertRaises(Exception):
ObjectExplorerService._generate_session_uri(param_set)
def test_list_databases_handles_query_failure(self):
"""Test that the list databases handler returns an error if the list databases query fails for any reason"""
# Set up the test with mock data
mock_query_results = [('database1',), ('database2',)]
connection_uri = 'someuri'
mock_cursor = MockCursor(mock_query_results)
mock_cursor.fetchall.side_effect = psycopg2.ProgrammingError('')
mock_connection = MockConnection(
dsn_parameters={
'host': 'myserver',
'dbname': 'postgres',
'user': '******'
},
cursor=mock_cursor)
mock_request_context = utils.MockRequestContext()
# 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
# Verify that calling the listdatabases handler returns the expected
# databases
params = ListDatabasesParams()
params.owner_uri = connection_uri
with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_connection)):
self.connection_service.handle_list_databases(mock_request_context, params)
self.assertIsNone(mock_request_context.last_notification_method)
self.assertIsNone(mock_request_context.last_notification_params)
self.assertIsNone(mock_request_context.last_response_params)
self.assertIsNotNone(mock_request_context.last_error_message)
def setUp(self):
"""Constructor"""
self.default_connection_key = 'server_db_user'
self.mock_connection_service = ConnectionService()
self.mock_server = JSONRPCServer(None, None)
self.mock_service_provider = ServiceProvider(self.mock_server, {},
None)
self.mock_service_provider._services[
constants.CONNECTION_SERVICE_NAME] = self.mock_connection_service
self.mock_service_provider._is_initialized = True
# Create connection information for use in the tests
self.connection_details = ConnectionDetails.from_data({})
self.connection_details.server_name = 'test_host'
self.connection_details.database_name = 'test_db'
self.connection_details.user_name = 'user'
self.expected_context_key = 'test_host|test_db|user'
self.expected_connection_uri = INTELLISENSE_URI + self.expected_context_key
self.test_uri = 'test_uri'
self.connection_info = ConnectionInfo(self.test_uri,
self.connection_details)
# Create mock CompletionRefresher to avoid calls to create separate thread
self.refresher_mock = mock.MagicMock()
self.refresh_method_mock = mock.MagicMock()
self.refresher_mock.refresh = self.refresh_method_mock
def test_build_connection_response(self):
"""Test that the connection response is built correctly"""
# Set up the test with mock data
server_name = 'testserver'
db_name = 'testdb'
user = '******'
mock_connection = MockConnection({
'host': server_name,
'dbname': db_name,
'user': user
})
connection_type = ConnectionType.EDIT
connection_details = ConnectionDetails.from_data(opts={})
owner_uri = 'test_uri'
connection_info = ConnectionInfo(owner_uri, connection_details)
connection_info._connection_map = {connection_type: mock_connection}
# If I build a connection response for the connection
response = pgsqltoolsservice.connection.connection_service._build_connection_response(
connection_info, connection_type)
# Then the response should have accurate information about the connection
self.assertEqual(response.owner_uri, owner_uri)
self.assertEqual(response.server_info.server_version, mock_connection.server_version)
self.assertEqual(response.server_info.is_cloud, False)
self.assertEqual(response.connection_summary.server_name, server_name)
self.assertEqual(response.connection_summary.database_name, db_name)
self.assertEqual(response.connection_summary.user_name, user)
self.assertEqual(response.type, connection_type)
def test_list_databases(self):
"""Test that the list databases handler correctly lists the connection's databases"""
# Set up the test with mock data
mock_query_results = [('database1',), ('database2',)]
connection_uri = 'someuri'
mock_connection = MockConnection(
dsn_parameters={
'host': 'myserver',
'dbname': 'postgres',
'user': '******'
},
cursor=MockCursor(mock_query_results))
mock_request_context = utils.MockRequestContext()
# 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
# Verify that calling the listdatabases handler returns the expected databases
params = ListDatabasesParams()
params.owner_uri = connection_uri
with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_connection)):
self.connection_service.handle_list_databases(mock_request_context, params)
expected_databases = [result[0] for result in mock_query_results]
self.assertEqual(mock_request_context.last_response_params.database_names, expected_databases)
def test_get_connection_creates_connection(self):
"""Test that get_connection creates a new connection when none 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
with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_connection)) as mock_psycopg2_connect:
# Open the connection
self.connection_service.connect(ConnectRequestParams(connection_details, connection_uri, connection_type))
# Get the connection
connection = self.connection_service.get_connection(connection_uri, connection_type)
self.assertEqual(connection, mock_connection)
mock_psycopg2_connect.assert_called_once()
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_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_connection = MockConnection(dsn_parameters={
'host': 'myserver',
'dbname': '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
# 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_connection)) as mock_psycopg2_connect:
response = self.connection_service.connect(params)
mock_psycopg2_connect.assert_not_called()
mock_connection.close.assert_not_called()
self.assertIsNotNone(response.connection_id)
def test_disconnect_for_invalid_connection(self):
"""Test that the disconnect method returns false when called on a connection that does not exist"""
# Set up the test with mock data
connection_uri = 'someuri'
connection_type_1 = ConnectionType.DEFAULT
mock_connection_1 = MockConnection(dsn_parameters={
'host': 'myserver1',
'dbname': 'postgres1',
'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)
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, ConnectionType.EDIT)
mock_connection_1.close.assert_not_called()
self.assertFalse(response)
def test_on_connect_sends_notification(self):
"""
Test that the service sends an intellisense ready notification after handling an on connect notification from the connection service.
This is a slightly more end-to-end test that verifies calling through to the queue layer
"""
# If: I create a new language service
service: LanguageService = self._init_service_with_flow_validator()
conn_info = ConnectionInfo(
'file://msuri.sql',
ConnectionDetails.from_data({
'host': None,
'dbname': 'TEST_DBNAME',
'user': '******'
}))
connect_result = mock.MagicMock()
connect_result.error_message = None
self.mock_connection_service.get_connection = mock.Mock(
return_value=mock.MagicMock())
self.mock_connection_service.connect = mock.MagicMock(
return_value=connect_result)
def validate_success_notification(response: IntelliSenseReadyParams):
self.assertEqual(response.owner_uri, conn_info.owner_uri)
# When: I notify of a connection complete for a given URI
self.flow_validator.add_expected_notification(
IntelliSenseReadyParams, INTELLISENSE_READY_NOTIFICATION,
validate_success_notification)
refresher_mock = mock.MagicMock()
refresh_method_mock = mock.MagicMock()
refresher_mock.refresh = refresh_method_mock
patch_path = 'pgsqltoolsservice.language.operations_queue.CompletionRefresher'
with mock.patch(patch_path) as refresher_patch:
refresher_patch.return_value = refresher_mock
task: threading.Thread = service.on_connect(conn_info)
# And when refresh is "complete"
refresh_method_mock.assert_called_once()
callback = refresh_method_mock.call_args[0][0]
self.assertIsNotNone(callback)
callback(None)
# Wait for task to return
task.join()
# Then:
# an intellisense ready notification should be sent for that URI
self.flow_validator.validate()
# ... and the scriptparseinfo should be created
info: ScriptParseInfo = service.get_script_parse_info(
conn_info.owner_uri)
self.assertIsNotNone(info)
# ... and the info should have the connection key set
self.assertEqual(info.connection_key,
OperationsQueue.create_key(conn_info))
def test_handle_close_session_incomplete_params(self):
# If: I close an OE session for with missing params
# NOTE: We only need to get the generate uri method to throw, we make sure it throws in all
# scenarios in a different test
rc = RequestFlowValidator().add_expected_error(
type(None), RequestFlowValidator.basic_error_validation)
params = ConnectionDetails.from_data({})
self.oe._handle_close_session_request(rc.request_context, params)
# Then:
# ... I should get an error response
rc.validate()
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 handle_change_database_request(self, request_context: RequestContext,
params: ChangeDatabaseRequestParams) -> bool:
"""change database of an existing connection or create a new connection
with default database from input"""
connection_info: ConnectionInfo = self.get_connection_info(params.owner_uri)
if connection_info is None:
return False
connection_info_params: Dict[str, str] = connection_info.details.options.copy()
connection_info_params["dbname"] = params.new_database
connection_details: ConnectionDetails = ConnectionDetails.from_data(connection_info_params)
connection_request_params: ConnectRequestParams = ConnectRequestParams(connection_details, params.owner_uri, ConnectionType.DEFAULT)
self.handle_connect_request(request_context, connection_request_params)
def _create_connection(self, session: ObjectExplorerSession, database_name: str) -> Optional[psycopg2.extensions.connection]:
conn_service = self._service_provider[utils.constants.CONNECTION_SERVICE_NAME]
options = session.connection_details.options.copy()
options['dbname'] = database_name
conn_details = ConnectionDetails.from_data(options)
key_uri = session.id + database_name
connect_request = ConnectRequestParams(conn_details, key_uri, ConnectionType.OBJECT_EXLPORER)
connect_result = conn_service.connect(connect_request)
if connect_result.error_message is not None:
raise RuntimeError(connect_result.error_message)
connection = conn_service.get_connection(key_uri, ConnectionType.OBJECT_EXLPORER)
return connection
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 = MockConnection(dsn_parameters={
'host': 'myserver',
'dbname': '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_handle_create_session_incomplete_params(self):
# Setup: Create an OE service
oe = ObjectExplorerService()
oe._service_provider = utils.get_mock_service_provider({})
# If: I create an OE session for with missing params
# NOTE: We only need to get the generate uri method to throw, we make sure it throws in all
# scenarios in a different test
rc = RequestFlowValidator().add_expected_error(
type(None), RequestFlowValidator.basic_error_validation)
params = ConnectionDetails.from_data({})
oe._handle_create_session_request(rc.request_context, params)
# Then:
# ... I should get an error response
rc.validate()
# ... A session should not have been created
self.assertDictEqual(oe._session_map, {})