def test_redirect(self): mocket = mock(websocket.WebSocket, strict=False) when(websocket).create_connection("ws://{}:{}".format( TestConnection.test_host, TestConnection.test_port)).thenReturn(mocket) proto_captor = captor(ANY(str)) when(mocket).send(proto_captor) msg = database_pb2.database_response() msg.redirect.leader_host = TestConnection.test_redirect_host msg.redirect.leader_port = TestConnection.test_redirect_port when(mocket).recv().thenReturn(msg.SerializeToString()) redirect_mocket = mock(websocket.WebSocket, strict=False) when(websocket).create_connection("ws://{}:{}".format( TestConnection.test_redirect_host, TestConnection.test_redirect_port)).thenReturn(redirect_mocket) msg = database_pb2.database_response() msg.resp.value = TestConnection.test_value when(redirect_mocket).recv().thenReturn(msg.SerializeToString()) self.assertEquals(self.connection.read(TestConnection.test_key), TestConnection.test_value) self.assertIsNotNone(proto_captor.value) send_json = json.loads(proto_captor.value) self.assertIn('bzn-api', send_json) self.assertEquals(send_json['bzn-api'], 'database') msg = bluzelle_pb2.bzn_msg() msg.ParseFromString(base64.b64decode(send_json['msg'])) self.assertEquals(msg.db.read.key, TestConnection.test_key) unstub()
def _send_request(self, host, port, msg): ws = websocket.create_connection("ws://{}:{}".format(host, port)) msg.db.header.db_uuid = self._uuid msg.db.header.transaction_id = random.randint(1, sys.maxint) req = {} req["bzn-api"] = "database" req["msg"] = base64.b64encode(msg.SerializeToString()) logger.debug("Sending: {}".format(msg)) logger.debug("-" * 60 + '\n') var = json.dumps(req) ws.send(json.dumps(req)) resp = database_pb2.database_response() resp.ParseFromString(ws.recv()) if resp.WhichOneof('success') == 'redirect': host = resp.redirect.leader_host port = resp.redirect.leader_port logger.debug("redirecting to leader at {}:{}...\n".format(host, port)) resp = self._send_request(host, port, msg) else: logger.debug("Response: \n{}".format(resp)) logger.debug("-" * 60 + '\n') ws.close() return resp
def test_keys(self): mocket = mock(websocket.WebSocket, strict=False) when(websocket).create_connection("ws://{}:{}".format(TestConnection.test_host, TestConnection.test_port)).thenReturn(mocket) proto_captor = captor(ANY(str)) when(mocket).send(proto_captor) keys = [TestConnection.test_key, TestConnection.test_key_other] msg = database_pb2.database_response() msg.keys.keys.extend(keys) when(mocket).recv().thenReturn(msg.SerializeToString()) test_connection = self.createTestConnection() self.assertListEqual(list(test_connection.keys()), keys) self.assertIsNotNone(proto_captor.value) send_json = json.loads(proto_captor.value) self.assertIn('bzn-api', send_json) self.assertEquals(send_json['bzn-api'], 'database') msg = bluzelle_pb2.bzn_msg() msg.ParseFromString(base64.b64decode(send_json['msg'])) unstub()
def test_create(self): mocket = mock(websocket.WebSocket, strict=False) when(websocket).create_connection("ws://{}:{}".format(TestConnection.test_host, TestConnection.test_port)).thenReturn(mocket) proto_captor = captor(ANY(str)) when(mocket).send(proto_captor) msg = database_pb2.database_response() when(mocket).recv().thenReturn(msg.SerializeToString()) test_connection = self.createTestConnection() self.assertIsNone(test_connection.create(TestConnection.test_key, TestConnection.test_value)) self.assertIsNotNone(proto_captor.value) send_json = json.loads(proto_captor.value) self.assertIn('bzn-api', send_json) self.assertEquals(send_json['bzn-api'], 'database') msg = bluzelle_pb2.bzn_msg() msg.ParseFromString(base64.b64decode(send_json['msg'])) self.assertEquals(msg.db.create.key, TestConnection.test_key) self.assertEquals(msg.db.create.value, TestConnection.test_value) self.assertEquals(msg.db.header.db_uuid, TestConnection.test_uuid) self.assertIsNotNone(msg.db.header.transaction_id) unstub()
def read_response(self, msg=None): resp = database_pb2.database_response() resp.ParseFromString(self._ws.recv()) if resp.WhichOneof('response') == 'redirect': host = resp.redirect.leader_host port = resp.redirect.leader_port self._connect(host, port) logger.debug("redirecting to leader at {}:{}.\n".format(host, port)) resp = self._send_request_sync(msg) else: logger.debug("Response: \n{}".format(resp)) return resp