diff --git a/aiogremlin/driver/protocol.py b/aiogremlin/driver/protocol.py index 02a3e905c2d35c98f7b50a8d3c38b68e695c35c2..5a53ef56e57ca69406c3d92203a12dd8f0701880 100644 --- a/aiogremlin/driver/protocol.py +++ b/aiogremlin/driver/protocol.py @@ -42,23 +42,15 @@ class GremlinServerWSProtocol(protocol.AbstractBaseProtocol): await func async def data_received(self, data, results_dict): - serializer_version = self._message_serializer.version data = data.decode('utf-8') - message = json.loads(data) + message = self._message_serializer.deserialize_message(json.loads(data)) request_id = message['requestId'] status_code = message['status']['code'] data = message['result']['data'] msg = message['status']['message'] if request_id in results_dict: result_set = results_dict[request_id] - if serializer_version == b"application/vnd.gremlin-v2.0+json": - aggregate_to = data['result']['meta'].get('aggregateTo', 'list') - else: - meta_aggregate_to = message['result']['meta']['@value'] - if len(meta_aggregate_to) > 1: - aggregate_to = meta_aggregate_to[1] - else: - aggregate_to = 'list' + aggregate_to = message['result']['meta'].get('aggregateTo', 'list') result_set.aggregate_to = aggregate_to if status_code == 407: @@ -72,18 +64,11 @@ class GremlinServerWSProtocol(protocol.AbstractBaseProtocol): result_set.queue_result(None) else: if data: - if serializer_version == b"application/vnd.gremlin-v2.0+json": - for result in data: - result = self._message_serializer.deserialize_message(result) - message = Message(status_code, result, msg) - result_set.queue_result(message) - else: - results = self._message_serializer.deserialize_message(data['@value']) - for result in results: - message = Message(status_code, result, msg) - result_set.queue_result(message) + for result in data: + result = self._message_serializer.deserialize_message(result) + message = Message(status_code, result, msg) + result_set.queue_result(message) else: - data = self._message_serializer.deserialize_message(data) message = Message(status_code, data, msg) result_set.queue_result(message) if status_code != 206: diff --git a/tests/test_gremlin_python/driver/test_driver_remote_connection.py b/tests/test_gremlin_python/driver/test_driver_remote_connection.py index 2cdbe84d78da733c5e92684c8427f8dfe995e758..883d4c11ff2e1e6cc0b4160f91401d4c17ac176e 100644 --- a/tests/test_gremlin_python/driver/test_driver_remote_connection.py +++ b/tests/test_gremlin_python/driver/test_driver_remote_connection.py @@ -38,6 +38,7 @@ class TestDriverRemoteConnection(object): statics.load_statics(globals()) g = Graph().traversal().withRemote(remote_connection) result = await g.V().limit(1).toList() + await remote_connection.close() @pytest.mark.asyncio async def test_traversals(self, remote_connection):