def test_send_message(self):
with io.BytesIO(b'') as stream:
# If:
# ... I create a JSON RPC writer
writer = JSONRPCWriter(stream, logger=utils.get_mock_logger())
# ... and I send a message
message = JSONRPCMessage.create_request('123', 'test/test', {})
writer.send_message(message)
# Then:
# ... The content-length header should be present
stream.seek(0)
header = stream.readline().decode('ascii')
self.assertRegex(
header,
re.compile('^Content-Length: [0-9]+\r\n$', re.IGNORECASE))
# ... There should be a blank line to signify the end of the headers
blank_line = stream.readline().decode('ascii')
self.assertEqual(blank_line, '\r\n')
# ... The JSON message as a dictionary should match the dictionary of the message
message_str = str.join(
os.linesep, [x.decode('UTF-8') for x in stream.readlines()])
message_dict = json.loads(message_str)
self.assertDictEqual(message_dict, message.dictionary)
def test_server_init(self):
# Setup: Create objects to init the server with
input_stream = io.BytesIO()
output_stream = io.BytesIO()
logger = utils.get_mock_logger()
# If: I create a server
server = JSONRPCServer(input_stream, output_stream, logger=logger)
# Then: The state should be initialized as defined
self.assertIsInstance(server.writer, JSONRPCWriter)
self.assertIsInstance(server.reader, JSONRPCReader)
self.assertIs(server._logger, logger)
self.assertEqual(server._version, '0')
self.assertFalse(server._stop_requested)
# ... The output queue should be empty
self.assertIsInstance(server._output_queue, Queue)
self.assertTrue(server._output_queue.all_tasks_done)
self.assertDictEqual(server._notification_handlers, {})
self.assertListEqual(server._shutdown_handlers, [])
# ... The threads shouldn't be assigned yet
self.assertIsNone(server._output_consumer)
self.assertIsNone(server._input_consumer)
# ... The built-in handlers should be assigned
self.assertTrue('echo' in server._request_handlers)
self.assertIsNotNone(server._request_handlers['echo'])
self.assertTrue('version' in server._request_handlers)
self.assertIsNotNone(server._request_handlers['version'].handler)
self.assertTrue('shutdown' in server._request_handlers)
self.assertIsNotNone(server._request_handlers['shutdown'].handler)
self.assertTrue('exit' in server._request_handlers)
self.assertIsNotNone(server._request_handlers['exit'].handler)
def test_handle_text_notification_none(self):
# Setup:
# ... Create a workspace service with mock callbacks and a workspace that always returns None
ws: WorkspaceService = WorkspaceService()
ws._logger = utils.get_mock_logger()
ws._text_change_callbacks = [MagicMock()]
ws._text_open_callbacks = [MagicMock()]
ws._text_close_callbacks = [MagicMock()]
ws._workspace, sf = self._get_mock_workspace(all_none=True)
nc: NotificationContext = utils.get_mock_notification_context()
# ... Create a list of methods call and parameters to call them with
test_calls = [
(ws._handle_did_change_text_doc,
self._get_change_text_doc_params(), ws._text_change_callbacks[0]),
(ws._handle_did_open_text_doc, self._get_open_text_doc_params(),
ws._text_open_callbacks[0]),
(ws._handle_did_close_text_doc, self._get_close_text_doc_params(),
ws._text_close_callbacks[0])
]
for call in test_calls:
# If: The workspace service receives a request to handle a file that shouldn't be processed
call[0](nc, call[1])
# Then: The associated notification callback should not have been called
call[2].assert_not_called()
def test_register(self):
"""Test registration of the service"""
# Setup:
# ... Create a mock service provider
server: JSONRPCServer = JSONRPCServer(None, None)
server.set_notification_handler = mock.MagicMock()
server.set_request_handler = mock.MagicMock()
provider: ServiceProvider = ServiceProvider(
server, {constants.CONNECTION_SERVICE_NAME: ConnectionService},
PG_PROVIDER_NAME, utils.get_mock_logger())
provider._is_initialized = True
conn_service: ConnectionService = provider[
constants.CONNECTION_SERVICE_NAME]
self.assertEqual(0, len(conn_service._on_connect_callbacks))
# If: I register a language service
service: LanguageService = LanguageService()
service.register(provider)
# Then:
# ... The notifications should have been registered
server.set_notification_handler.assert_called()
server.set_request_handler.assert_called()
self.assertEqual(1, len(conn_service._on_connect_callbacks))
self.assertEqual(1, server.count_shutdown_handlers())
# ... The service provider should have been stored
self.assertIs(service._service_provider, provider) # noqa
def test_read_next_chunk_resize(self):
# Setup: Create a byte array for test input
test_bytes = bytearray(b'1234567890')
with io.BytesIO(test_bytes) as stream:
# If:
# ... I create a reader with an artificially low initial buffer size
# and prefill the buffer
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
reader._buffer = bytearray(5)
reader._buffer_end_offset = 4
# ... and I read a chunk from the stream
result = reader._read_next_chunk()
# Then:
# ... The read should have succeeded
self.assertTrue(result, True)
# ... The size of the buffer should have doubled
self.assertEqual(len(reader._buffer), 10)
# ... The buffer end offset should be the size of the buffer
self.assertEqual(reader._buffer_end_offset, 10)
# ... The buffer should contain the first 6 elements of the test data
expected = test_bytes[:6]
actual = reader._buffer[4:]
self.assertEqual(actual, expected)
def test_read_multiple_messages(self):
# Setup:
# ... Create an input stream with two messages
test_bytes = b'Content-Length: 30\r\n\r\n{"method":"test", "params":{}}'
input_stream = io.BytesIO(test_bytes + test_bytes)
output_stream = io.BytesIO()
# ... Create a server that uses the input and output streams
server = JSONRPCServer(input_stream, output_stream, logger=utils.get_mock_logger())
# ... Patch the server to not dispatch a message
dispatch_mock = mock.MagicMock()
server._dispatch_message = dispatch_mock
# If: I start the server, run it for a bit, and stop it
server.start()
time.sleep(1)
server.stop()
server.wait_for_exit()
# Then: The dispatch method should have been called twice
expected_output = JSONRPCMessage.from_dictionary({"method": "test", "params": {}})
self.assertEqual(len(dispatch_mock.mock_calls), 2)
self.assertDictEqual(dispatch_mock.mock_calls[0][1][0].dictionary, expected_output.dictionary)
self.assertDictEqual(dispatch_mock.mock_calls[1][1][0].dictionary, expected_output.dictionary)
# Teardown: All background threads should be shut down.
self.assertFalse(server._input_consumer.isAlive())
self.assertFalse(server._output_consumer.isAlive())
def _get_service_provider(services: int = 1) -> ServiceProvider:
# If: I create a new service provider
server = JSONRPCServer(None, None)
logger = utils.get_mock_logger()
services = {'service_name' + str(x): TestServiceProvider._TestService for x in range(0, services)}
sp = ServiceProvider(server, services, logger)
return sp
def test_read_next_chunk_eof(self):
with io.BytesIO() as stream:
# If:
# ... I create a reader with a stream that has no bytes
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
# ... and I read a chunk from the stream
# Then: I should get an exception
with self.assertRaises(EOFError):
reader._read_next_chunk()
def test_closes(self):
with io.BytesIO(b'') as stream:
# If:
# ... I create a JSON RPC reader with an opened stream
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
# ... and I close the reader
reader.close()
# Then: The stream should be closed
self.assertTrue(reader.stream.closed)
def test_create_standard_encoding(self):
with io.BytesIO(b'') as stream:
# If: I create a JSON RPC reader with a stream without specifying the encoding
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
# Then: The stream and encoding should be set appropriately
self.assertIsNotNone(reader)
self.assertIs(reader.stream, stream)
self.assertEqual(reader.encoding, 'UTF-8')
self.assertEqual(reader._read_state,
JSONRPCReader.ReadState.Header)
def test_close(self):
with io.BytesIO(b'123') as stream:
# If:
# ... I create a JSON RPC writer with an opened stream
writer = JSONRPCWriter(stream, logger=utils.get_mock_logger())
# ... and I close the writer
writer.close()
# Then: The stream should be closed
self.assertTrue(writer.stream.closed)
def test_dispatch_notification_no_handler():
# If: I dispatch a message that has no handler
logger = utils.get_mock_logger()
message = JSONRPCMessage.create_notification('non_existent', {})
server = JSONRPCServer(None, None, logger=logger)
server._dispatch_message(message)
# Then:
# ... Nothing should have happened
# TODO: Capture that an error was sent
# ... A warning should have been logged
logger.warn.assert_called_once()
def test_read_next_chunk_closed(self):
# Setup: Create a stream that has already been closed
stream = io.BytesIO()
stream.close()
# If:
# ... I create a reader with a closed stream
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
# ... and I read a chunk from the stream
# Then: I should get an exception
with self.assertRaises(ValueError):
reader._read_next_chunk()
def test_create_nonstandard_encoding(self):
with io.BytesIO(b'') as stream:
# If: I create a JSON RPC reader with a non-standard encoding
reader = JSONRPCReader(stream,
encoding="ascii",
logger=utils.get_mock_logger())
# Then: The stream and encoding should be set appropriately
self.assertIsNotNone(reader)
self.assertIs(reader.stream, stream)
self.assertEqual(reader.encoding, 'ascii')
self.assertEqual(reader._read_state,
JSONRPCReader.ReadState.Header)
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_read_message_invalid_json(self):
# Setup: Reader with a stream that has an invalid message
test_bytes = bytearray(b'Content-Length: 10\r\n\r\nabcdefghij')
with io.BytesIO(test_bytes) as stream:
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
reader._buffer = bytearray(100)
# If: I read a message
# Then:
# ... It should throw an exception
with self.assertRaises(ValueError):
reader.read_message()
# ... The buffer should be trashed
self.assertEqual(len(reader._buffer), reader.DEFAULT_BUFFER_SIZE)
def test_closes_exception():
# Setup: Patch the stream to have a custom close handler
stream = io.BytesIO(b'')
close_orig = stream.close
stream.close = mock.MagicMock(side_effect=AttributeError)
# If: Close a reader and it throws an exception
logger = utils.get_mock_logger()
reader = JSONRPCReader(stream, logger=logger)
reader.close()
# Then: There should not have been an exception throws
logger.exception.assert_called_once()
# Cleanup: Close the stream
close_orig()
def test_read_headers_not_found(self):
# Setup: Create a reader with a loaded buffer that does not contain the \r\n\r\n control
reader = JSONRPCReader(None, logger=utils.get_mock_logger())
reader._buffer = bytearray(b'1234567890')
reader._buffer_end_offset = len(reader._buffer)
# If: I look for a header block in the buffer
result = reader._try_read_headers()
# Then:
# ... I should not have found any
self.assertFalse(result)
# ... The current reading position of the buffer should not have moved
self.assertEqual(reader._read_offset, 0)
self.assertEqual(reader._read_state, JSONRPCReader.ReadState.Header)
def test_handle_did_change_config(self):
# Setup: Create a workspace service with two mock config change handlers
ws: WorkspaceService = WorkspaceService()
ws._logger = utils.get_mock_logger()
ws._config_change_callbacks = [MagicMock(), MagicMock()]
# If: The workspace receives a config change notification
nc: NotificationContext = utils.get_mock_notification_context()
params: DidChangeConfigurationParams = DidChangeConfigurationParams.from_dict(
{
'settings': {
'sql': {
'intellisense': {
'enable_intellisense': False
}
},
'pgsql': {
'format': {
'keyword_case': 'upper',
'identifier_case': 'lower',
'strip_comments': True,
'reindent': False,
}
}
}
})
ws._handle_did_change_config(nc, params)
# Then:
# ... No notifications should have been sent
nc.send_notification.assert_not_called()
# ... The config should have been updated
self.assertIs(ws.configuration, params.settings)
self.assertEqual(ws.configuration.pgsql.format.keyword_case, 'upper')
self.assertEqual(ws.configuration.pgsql.format.identifier_case,
'lower')
self.assertTrue(ws.configuration.pgsql.format.strip_comments)
self.assertFalse(ws.configuration.pgsql.format.reindent)
# ... And default values that weren't specified in the notification are preserved
self.assertTrue(ws.configuration.sql.intellisense.enable_suggestions)
# ... The mock config change callbacks should have been called
for callback in ws._config_change_callbacks:
callback.assert_called_once_with(params.settings)
def test_read_recover_from_content_message(self):
test_string = b'Content-Length: 10\r\n\r\nabcdefghij' + \
b'Content-Length: 32\r\n\r\n{"method":"test", "params":null}'
test_bytes = bytearray(test_string)
with io.BytesIO(test_bytes) as stream:
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
reader._buffer = bytearray(100)
# If: I read a message with invalid content
# Then: I should get an exception
with self.assertRaises(ValueError):
reader.read_message()
# If: I read another valid message
msg = reader.read_message()
# Then: I should have a valid message
self.assertIsNotNone(msg)
def test_handle_text_notification_success(self):
# Setup:
# ... Create a workspace service with a mock callback and a workspace that returns a mock script file
ws: WorkspaceService = WorkspaceService()
ws._logger = utils.get_mock_logger()
ws._workspace, sf = self._get_mock_workspace(False)
ws._text_change_callbacks = [MagicMock()]
ws._text_open_callbacks = [MagicMock()]
ws._text_close_callbacks = [MagicMock()]
# ... Create a mock notification context
nc: NotificationContext = utils.get_mock_notification_context()
# ... Create a list of method calls and parameters to call them with
test_calls = [
(ws._handle_did_change_text_doc, ws._text_change_callbacks[0],
self._get_change_text_doc_params(),
self._test_handle_text_change_helper),
(ws._handle_did_open_text_doc, ws._text_open_callbacks[0],
self._get_open_text_doc_params(), None),
(ws._handle_did_close_text_doc, ws._text_close_callbacks[0],
self._get_close_text_doc_params(), None)
]
for call in test_calls:
# If: The workspace service receives a notification
call[0](nc, call[2])
# Then:
# ... The callback should have been called with the script file
call[1].assert_called_once_with(sf)
# ... The notification sender should not have not been called
nc.send_notification.assert_not_called()
# ... Any additional validation should pass
if call[3] is not None:
call[3](call[2], sf)
# ... Get, Open, and Close file should all have been called
ws._workspace.get_file.assert_called_once()
ws._workspace.open_file.assert_called_once()
ws._workspace.close_file.assert_called_once()
def test_register(self):
# Setup:
# ... Create a mock service provider
server: JSONRPCServer = JSONRPCServer(None, None)
server.set_notification_handler = mock.MagicMock()
server.set_request_handler = mock.MagicMock()
sp: ServiceProvider = ServiceProvider(server, {},
utils.get_mock_logger())
# If: I register a OE service
oe = ObjectExplorerService()
oe.register(sp)
# Then:
# ... The service should have registered its request handlers
server.set_request_handler.assert_called()
server.set_notification_handler.assert_not_called()
# ... The service provider should have been stored
self.assertIs(oe._service_provider, sp)
def test_read_next_chunk_success(self):
# Setup: Create a byte array for test input
test_bytes = bytearray(b'123')
with io.BytesIO(test_bytes) as stream:
# If: I attempt to read a chunk from the stream
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
result = reader._read_next_chunk()
# Then:
# ... The result should be true
self.assertTrue(result)
# ... The buffer should contain 3 byte and should not have been read yet
self.assertEqual(reader._buffer_end_offset, len(test_bytes))
self.assertEqual(reader._read_offset, 0)
# ... The buffer should now contain the bytes from the stream
buffer_contents = reader._buffer[:len(test_bytes)]
self.assertEqual(buffer_contents, test_bytes)
def test_register(self):
# Setup:
# ... Create a mock service provider
server: JSONRPCServer = JSONRPCServer(None, None)
server.set_notification_handler = MagicMock()
server.set_request_handler = MagicMock()
sp: ServiceProvider = ServiceProvider(server, {},
utils.get_mock_logger())
# If: I register a workspace service
ws: WorkspaceService = WorkspaceService()
ws.register(sp)
# Then:
# ... The notifications should have been registered
server.set_notification_handler.assert_called()
server.set_request_handler.assert_not_called()
# ... The service provider should have been stored
self.assertIs(ws._service_provider, sp)
def test_init(self):
# If: I create a new service provider
server = JSONRPCServer(None, None)
logger = utils.get_mock_logger()
mock_service = mock.MagicMock(return_value={})
services = {'service_name': mock_service}
sp = ServiceProvider(server, services, logger)
# Then:
# ... The properties should return the values I set (server/logger)
self.assertFalse(sp._is_initialized)
self.assertIs(sp._server, server)
self.assertIs(sp.server, server)
self.assertIs(sp._logger, logger)
self.assertIs(sp.logger, logger)
# ... The services should be transformed and called
self.assertIsInstance(sp._services, dict)
self.assertTrue('service_name' in sp._services)
mock_service.assert_called_once()
def test_dispatch_notification_normal(self):
# Setup: Create a server with a single handler that has none for the deserialization class
config = IncomingMessageConfiguration('test/test', _TestParams)
handler = mock.MagicMock()
server = JSONRPCServer(None, None, logger=utils.get_mock_logger())
server.set_notification_handler(config, handler)
# If: I dispatch a message that has none set for the deserialization class
params = {}
message = JSONRPCMessage.create_notification('test/test', params)
server._dispatch_message(message)
# Then:
# ... The handler should have been called
handler.assert_called_once()
# ... The parameters to the handler should have been a request context and params
self.assertIsInstance(handler.mock_calls[0][1][0], NotificationContext)
self.assertIs(handler.mock_calls[0][1][0]._queue, server._output_queue)
self.assertIsInstance(handler.mock_calls[0][1][1], _TestParams)
def test_registration(self):
# Setup:
# ... Create a mock service provider
server: JSONRPCServer = JSONRPCServer(None, None)
server.set_notification_handler = mock.MagicMock()
server.set_request_handler = mock.MagicMock()
sp: ServiceProvider = ServiceProvider(server, {}, PG_PROVIDER_NAME,
utils.get_mock_logger())
# If: I register a scripting service
ss: ScriptingService = ScriptingService()
ss.register(sp)
# Then:
# ... The service should have registered its request handlers
server.set_request_handler.assert_called()
server.set_notification_handler.assert_not_called()
# ... The service provider should have been stored
self.assertIs(ss._service_provider, sp)
def test_read_headers_no_colon(self):
# Setup: Create a reader with a buffer that contains the control sequence but does not
# match the header format
test_buffer = bytearray(b'1234567890\r\n\r\n')
reader = JSONRPCReader(None, logger=utils.get_mock_logger())
reader._buffer = test_buffer
reader._buffer_end_offset = len(reader._buffer)
reader._read_offset = 1
# If: I look for a header block in the buffer
# Then:
# ... I should get an exception b/c of the malformed header
with self.assertRaises(KeyError):
reader._try_read_headers()
# ... The current reading position of the buffer should be reset to 0
self.assertEqual(reader._read_offset, 0)
# ... The buffer should have been trashed
self.assertIsNot(reader._buffer, test_buffer)
def test_read_headers_success(self):
# Setup: Create a reader with a loaded buffer that contains a a complete header
reader = JSONRPCReader(None, logger=utils.get_mock_logger())
reader._buffer = bytearray(b'Content-Length: 56\r\n\r\n')
reader._buffer_end_offset = len(reader._buffer)
# If: I look for a header block in the buffer
result = reader._try_read_headers()
# Then:
# ... I should have found a header
self.assertTrue(result)
# ... The current reading position should have moved to the end of the buffer
self.assertEqual(reader._read_offset, len(reader._buffer))
self.assertEqual(reader._read_state, JSONRPCReader.ReadState.Content)
# ... The headers should have been stored
self.assertEqual(reader._expected_content_length, 56)
self.assertDictEqual(reader._headers, {'content-length': '56'})
def test_read_multiple_messages(self):
test_string = b'Content-Length: 32\r\n\r\n{"method":"test", "params":null}'
test_bytes = bytearray(test_string + test_string)
with io.BytesIO(test_bytes) as stream:
reader = JSONRPCReader(stream, logger=utils.get_mock_logger())
reader._buffer = bytearray(100)
# If:
# ... I read a message
msg1 = reader.read_message()
# ... And I read another message
msg2 = reader.read_message()
# Then:
# ... The messages should be real
self.assertIsNotNone(msg1)
self.assertIsNotNone(msg2)
# ... The buffer should have been trashed
self.assertEqual(len(reader._buffer), reader.DEFAULT_BUFFER_SIZE)
