def test_edit_initialize(self):
request_context = utils.MockRequestContext()
params = None
edit_session = mock.MagicMock()
edit_session._edit_initialize(request_context, params)
edit_session._edit_initialize.assert_called()
def test_commit_when_edit_session_is_not_available(self):
request_context = utils.MockRequestContext()
request = EditCommitRequest()
request.owner_uri = 'owner_uri'
edit_session = mock.MagicMock()
self._service_under_test._active_sessions[
request.owner_uri] = edit_session
self._service_under_test._edit_commit(request_context, request)
args = edit_session.commit_edit.call_args[0]
success_callback = args[1]
success_callback()
request_context.send_response.assert_called_once()
failure_callback = args[2]
error_message = 'Test error'
failure_callback(error_message)
self.assertEqual(error_message, request_context.last_error_message)
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 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_dmp_capabilities_have_backup_options(self):
"""Test that the capabilities returned for a DMP capabilities request include backup options"""
# Setup: Create a request context with mocked out send_* methods and set up the capabilities service
rc = utils.MockRequestContext()
capabilities_service = CapabilitiesService()
workspace_service = WorkspaceService()
capabilities_service._service_provider = utils.get_mock_service_provider(
{constants.WORKSPACE_SERVICE_NAME: workspace_service})
# If: I request the dmp capabilities of this server
capabilities_service._handle_dmp_capabilities_request(rc, None)
# Then: The response should include backup capabilities
rc.send_response.assert_called_once()
capabilities_result = rc.send_response.mock_calls[0][1][0]
features = capabilities_result.capabilities.features
backup_options_list = [
feature for feature in features if feature.feature_name == 'backup'
]
# There should be exactly one feature containing backup options
self.assertEqual(len(backup_options_list), 1)
backup_options = backup_options_list[0]
# The backup options should be enabled
self.assertTrue(backup_options.enabled)
# And the backup options should contain at least 1 option
self.assertGreater(len(backup_options.options_metadata), 0)
def test_handle_database_change_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)
self.connection_service.get_connection_info = mock.MagicMock()
params: ChangeDatabaseRequestParams = ChangeDatabaseRequestParams.from_dict({
'owner_uri': 'someUri',
'new_database': 'newDb'
})
self.connection_service.handle_change_database_request(rc, params)
connection_thread = self.connection_service.owner_to_thread_map[params.owner_uri]
self.assertIsNotNone(connection_thread)
connection_thread.join()
self.connection_service.connect.assert_called_once()
# ... 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_default_completion_items(self):
"""
Test that the completion handler returns a set of default values
when not connected to any URI
"""
# If: The script file exists
input_text = 'create tab'
doc_position = TextDocumentPosition.from_dict({
'text_document': {
'uri': self.default_uri
},
'position': {
'line': 0,
'character': 10 # end of 'tab' word
}
})
context: RequestContext = utils.MockRequestContext()
config = Configuration()
config.sql.intellisense.enable_intellisense = True
self.mock_workspace_service._configuration = config
workspace, script_file = self._get_test_workspace(True, input_text)
self.mock_workspace_service._workspace = workspace
service: LanguageService = self._init_service()
service._valid_uri.add(doc_position.text_document.uri)
# When: I request completion item
service.handle_completion_request(context, doc_position)
# Then:
# ... An default completion set should be sent over the notification
context.send_response.assert_called_once()
completions: List[CompletionItem] = context.last_response_params
self.assertTrue(len(completions) > 0)
self.verify_match('TABLE', completions, Range.from_data(0, 7, 0, 10))
def test_format_doc_no_pgsql_format(self):
"""
Test that the format codepath succeeds even if the configuration options aren't defined
"""
input_text = 'select * from foo where id in (select id from bar);'
context: RequestContext = utils.MockRequestContext()
self.mock_workspace_service._configuration = None
workspace, script_file = self._get_test_workspace(True, input_text)
self.mock_workspace_service._workspace = workspace
service: LanguageService = self._init_service()
format_options = FormattingOptions()
format_options.insert_spaces = False
format_params = DocumentFormattingParams()
format_params.options = format_options
format_params.text_document = self.default_text_document_id
# add uri to valid uri set ensure request passes uri check
# normally done in flavor change handler, but we are not testing that here
service._valid_uri.add(format_params.text_document.uri)
# When: I have no useful formatting defaults defined
service.handle_doc_format_request(context, format_params)
# Then:
# ... There should be no changes to the doc
context.send_response.assert_called_once()
edits: List[TextEdit] = context.last_response_params
self.assertTrue(len(edits) > 0)
self.assert_range_equals(edits[0].range,
Range.from_data(0, 0, 0, len(input_text)))
self.assertEqual(edits[0].new_text, input_text)
def setUp(self):
"""Set up mock objects for testing the query execution service.
Ran before each unit test.
"""
# set up mock connection
self.rows = [(1, 'Text 1'), (2, 'Text 2')]
self.cursor = utils.MockCursor(self.rows)
self.mock_pymysql_connection = utils.MockPyMySQLConnection(
parameters={
'host': 'test',
'dbname': 'test',
})
self.connection = MockMySQLServerConnection()
self.connection.cursor.return_value = self.cursor
self.cursor.connection = self.connection
self.connection_service = ConnectionService()
self.request_context = utils.MockRequestContext()
# setup mock query_execution_service
self.query_execution_service = QueryExecutionService()
self.service_provider = ServiceProvider(None, {},
constants.MYSQL_PROVIDER_NAME)
self.service_provider._services = {
constants.CONNECTION_SERVICE_NAME: self.connection_service
}
self.service_provider._is_initialized = True
self.query_execution_service._service_provider = self.service_provider
def connection_side_effect(owner_uri: str,
connection_type: ConnectionType):
return self.connection
self.connection_service.get_connection = mock.Mock(
side_effect=connection_side_effect)
def setUp(self):
"""Set up the data for tasks"""
self.mock_action = mock.Mock(return_value=TaskResult(TaskStatus.SUCCEEDED))
self.task_name = 'test_name'
self.task_description = 'test_description'
self.task_provider = constants.PROVIDER_NAME
self.server_name = 'test_server'
self.database_name = 'test_db'
self.request_context = utils.MockRequestContext()
def test_version_request():
# If: I send a request for the version
rc = utils.MockRequestContext()
server = JSONRPCServer(None, None)
server._handle_version_request(rc, None)
# Then: I should get a response
rc.send_response.assert_called_once_with(server._version)
rc.send_error.assert_not_called()
rc.send_notification.assert_not_called()
def test_echo_request():
# If: I send a request for an echo
rc = utils.MockRequestContext()
params = {}
JSONRPCServer._handle_echo_request(rc, params)
# Then: The params should have been echoed back
rc.send_response.assert_called_once_with(params)
rc.send_notification.assert_not_called()
rc.send_error.assert_not_called()
def test_update_cell_with_no_active_session(self):
update_cell_request = UpdateCellRequest()
update_cell_request.owner_uri = 'test_owner_uri'
request_context = utils.MockRequestContext()
with self.assertRaises(KeyError):
self._service_under_test._update_cell(request_context,
update_cell_request)
def test_dispose_when_edit_session_is_not_available(self):
request_context = utils.MockRequestContext()
request = DisposeRequest()
request.owner_uri = 'owner_uri'
self._service_under_test._dispose(request_context, request)
self.assertEqual(request_context.last_error_message,
'Edit data session not found')
def test_list_databases_handles_invalid_uri(self):
"""Test that the connection/listdatabases handler returns an error when the given URI is unknown"""
mock_request_context = utils.MockRequestContext()
params = ListDatabasesParams()
params.owner_uri = 'unknown_uri'
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 test_handle_definition_request_should_return_empty_if_query_file_do_not_exist(
self):
# If: The script file doesn't exist (there is an empty workspace)
context: RequestContext = utils.MockRequestContext()
self.mock_workspace_service._workspace = Workspace()
service: LanguageService = self._init_service()
service.handle_definition_request(context, self.default_text_position)
context.send_response.assert_called_once()
self.assertEqual(context.last_response_params, [])
def test_handle_cancellation_no_match(self):
"""Test that handling a cancellation request returns false if there is no matching connection to cancel"""
# Set up a mock request handler
request_context = utils.MockRequestContext()
# If I call the cancellation request handler
cancel_params = CancelConnectParams(self.owner_uri, self.connection_type)
self.connection_service.handle_cancellation_request(request_context, cancel_params)
# Then the handler should have responded false to indicate that no matching connection was in progress
request_context.send_response.assert_called_once_with(False)
def test_handle_definition_request_intellisense_off(self):
request_context: RequestContext = utils.MockRequestContext()
config = Configuration()
config.sql.intellisense.enable_intellisense = False
self.mock_workspace_service._configuration = config
language_service = self._init_service()
language_service.handle_definition_request(request_context,
self.default_text_position)
request_context.send_response.assert_called_once()
self.assertEqual(request_context.last_response_params, [])
def test_format_mysql_doc_range(self):
"""
Test that the format document range codepath works as expected with a mysql doc
"""
# set up service provider with mysql connection
self.mock_service_provider = ServiceProvider(self.mock_server, {},
MYSQL_PROVIDER_NAME, None)
self.mock_service_provider._services[
constants.WORKSPACE_SERVICE_NAME] = self.mock_workspace_service
self.mock_service_provider._services[
constants.CONNECTION_SERVICE_NAME] = self.mock_connection_service
self.mock_service_provider._is_initialized = True
# If: The script file doesn't exist (there is an empty workspace)
input_lines: List[str] = [
'select * from t1',
'select * from foo where id in (select id from bar);'
]
input_text = '\n'.join(input_lines)
expected_output = '\n'.join([
'SELECT *', 'FROM foo', 'WHERE id IN', '\t\t\t\t(SELECT id',
'\t\t\t\t\tFROM bar);'
])
context: RequestContext = utils.MockRequestContext()
config = Configuration()
config.my_sql = MySQLConfiguration()
config.my_sql.format.keyword_case = 'upper'
self.mock_workspace_service._configuration = config
workspace, script_file = self._get_test_workspace(True, input_text)
self.mock_workspace_service._workspace = workspace
service: LanguageService = self._init_service()
format_options = FormattingOptions()
format_options.insert_spaces = False
format_params = DocumentRangeFormattingParams()
format_params.options = format_options
format_params.text_document = self.default_text_document_id
# add uri to valid uri set ensure request passes uri check
# normally done in flavor change handler, but we are not testing that here
service._valid_uri.add(format_params.text_document.uri)
# When: I request format the 2nd line of a document
format_params.range = Range.from_data(1, 0, 1, len(input_lines[1]))
service.handle_doc_range_format_request(context, format_params)
# Then:
# ... only the 2nd line should be formatted
context.send_response.assert_called_once()
edits: List[TextEdit] = context.last_response_params
self.assertTrue(len(edits) > 0)
self.assert_range_equals(edits[0].range, format_params.range)
self.assertEqual(edits[0].new_text, expected_output)
def test_initialization_request(self):
# Setup: Create a request context with mocked out send_* methods
rc = utils.MockRequestContext()
# If: I request the language service capabilities of this server
CapabilitiesService._handle_initialize_request(rc, None)
# Then: A response should have been sent that is a Capabilities result
rc.send_notification.assert_not_called()
rc.send_error.assert_not_called()
rc.send_response.assert_called_once()
self.assertIsInstance(rc.send_response.mock_calls[0][1][0],
InitializeResult)
def test_dispose_when_edit_session_available(self):
request_context = utils.MockRequestContext()
edit_session = mock.MagicMock()
request = DisposeRequest()
request.owner_uri = 'owner_uri'
self._service_under_test._active_sessions[
request.owner_uri] = edit_session
self._service_under_test._dispose(request_context, request)
self.assertEqual(0, len(self._service_under_test._active_sessions))
def test_format_mysql_doc(self):
"""
Test that the format document codepath works as expected with a mysql doc
"""
# set up service provider with mysql connection
self.mock_service_provider = ServiceProvider(self.mock_server, {},
MYSQL_PROVIDER_NAME, None)
self.mock_service_provider._services[
constants.WORKSPACE_SERVICE_NAME] = self.mock_workspace_service
self.mock_service_provider._services[
constants.CONNECTION_SERVICE_NAME] = self.mock_connection_service
self.mock_service_provider._is_initialized = True
# If: We have a basic string to be formatted
input_text = 'select * from foo where id in (select id from bar);'
# Note: sqlparse always uses '\n\ for line separator even on windows.
# For now, respecting this behavior and leaving as-is
expected_output = '\n'.join([
'SELECT *', 'FROM foo', 'WHERE id IN', '\t\t\t\t(SELECT id',
'\t\t\t\t\tFROM bar);'
])
context: RequestContext = utils.MockRequestContext()
config = Configuration()
config.my_sql = MySQLConfiguration()
config.my_sql.format.keyword_case = 'upper'
self.mock_workspace_service._configuration = config
workspace, script_file = self._get_test_workspace(True, input_text)
self.mock_workspace_service._workspace = workspace
service: LanguageService = self._init_service()
format_options = FormattingOptions()
format_options.insert_spaces = False
format_params = DocumentFormattingParams()
format_params.options = format_options
format_params.text_document = self.default_text_document_id
# add uri to valid uri set ensure request passes uri check
# normally done in flavor change handler, but we are not testing that here
service._valid_uri.add(format_params.text_document.uri)
# When: I request document formatting
service.handle_doc_format_request(context, format_params)
# Then:
# ... The entire document text should be formatted
context.send_response.assert_called_once()
edits: List[TextEdit] = context.last_response_params
self.assertTrue(len(edits) > 0)
self.assert_range_equals(edits[0].range,
Range.from_data(0, 0, 0, len(input_text)))
self.assertEqual(edits[0].new_text, expected_output)
def test_handle_database_change_request_with_empty_connection_info_for_false(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)
params: ChangeDatabaseRequestParams = ChangeDatabaseRequestParams.from_dict({
'owner_uri': 'someUri',
'new_database': 'newDb'
})
result: bool = self.connection_service.handle_change_database_request(rc, params)
self.assertEqual(result, False)
def test_shutdown_request(self):
# If: I send a request for the service to shutdown
rc = utils.MockRequestContext()
handler = mock.MagicMock()
server = JSONRPCServer(None, None, logger=utils.get_mock_logger())
server.add_shutdown_handler(handler)
server._handle_shutdown_request(rc, None)
# Then:
# ... The server should be shutting down
self.assertTrue(server._stop_requested)
# ... The shutdown handler should be called
handler.assert_called_once()
def test_handle_disconnect_request_unknown_uri(self):
"""Test that the handle_disconnect_request method returns false when the given URI is unknown"""
# Setup: Create a mock request context
rc = utils.MockRequestContext()
# If: I request to disconnect an unknown URI
params: DisconnectRequestParams = DisconnectRequestParams.from_dict(
{'ownerUri': 'nonexistent'})
self.connection_service.handle_disconnect_request(rc, params)
# Then: Send result should have been called once with False
rc.send_response.assert_called_once_with(False)
rc.send_notification.assert_not_called()
rc.send_error.assert_not_called()
def test_handle_cancellation_request(self):
"""Test that handling a cancellation request modifies the cancellation token for a matched connection"""
# Set up the connection service with a mock request handler and cancellation token
cancellation_key = (self.owner_uri, self.connection_type)
cancellation_token = CancellationToken()
self.connection_service._cancellation_map[cancellation_key] = cancellation_token
request_context = utils.MockRequestContext()
# If I call the cancellation request handler
cancel_params = CancelConnectParams(self.owner_uri, self.connection_type)
self.connection_service.handle_cancellation_request(request_context, cancel_params)
# Then the handler should have responded and set the cancellation flag
request_context.send_response.assert_called_once_with(True)
self.assertTrue(cancellation_token.canceled)
def test_completion_file_not_found(self):
"""
Test that the completion handler returns empty if the intellisense
is disabled
"""
# If: The script file doesn't exist (there is an empty workspace)
context: RequestContext = utils.MockRequestContext()
self.mock_workspace_service._workspace = Workspace()
service: LanguageService = self._init_service()
# When: I request completion item
service.handle_completion_request(context, self.default_text_position)
# Then:
# ... An empty completion should be sent over the notification
context.send_response.assert_called_once()
self.assertEqual(context.last_response_params, [])
def test_dmp_capabilities_request(self):
# Setup: Create a request context with mocked out send_* methods and set up the capabilities service
rc = utils.MockRequestContext()
capabilities_service = CapabilitiesService()
workspace_service = WorkspaceService()
capabilities_service._service_provider = utils.get_mock_service_provider(
{constants.WORKSPACE_SERVICE_NAME: workspace_service})
# If: I request the dmp capabilities of this server
capabilities_service._handle_dmp_capabilities_request(rc, None)
# Then: A response should have been sent that is a Capabilities result
rc.send_notification.assert_not_called()
rc.send_error.assert_not_called()
rc.send_response.assert_called_once()
self.assertIsInstance(rc.send_response.mock_calls[0][1][0],
CapabilitiesResult)
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_completion_intellisense_off(self):
"""
Test that the completion handler returns empty if the intellisense
is disabled
"""
# If: intellisense is disabled
context: RequestContext = utils.MockRequestContext()
config = Configuration()
config.sql.intellisense.enable_intellisense = False
self.mock_workspace_service._configuration = config
service: LanguageService = self._init_service()
# When: I request completion item
service.handle_completion_request(context, self.default_text_position)
# Then:
# ... An empty completion should be sent over the notification
context.send_response.assert_called_once()
self.assertEqual(context.last_response_params, [])