def test_set_exception(): status = StatusMessage() exc = Exception('exception code') status.set_exception(exc) assert status.proto.code == jina_pb2.StatusProto.ERROR assert status.proto.description == repr(exc) assert status.proto.exception.name == exc.__class__.__name__
async def dry_run(self, empty, context) -> jina_pb2.StatusProto: """ Process the the call requested by having a dry run call to every Executor in the graph :param empty: The service expects an empty protobuf message :param context: grpc context :returns: the response request """ from docarray import DocumentArray from jina.clients.request import request_generator from jina.enums import DataInputType from jina.serve.executors import __dry_run_endpoint__ da = DocumentArray() try: req_iterator = request_generator( exec_endpoint=__dry_run_endpoint__, data=da, data_type=DataInputType.DOCUMENT, ) async for _ in self.streamer.stream(request_iterator=req_iterator): pass status_message = StatusMessage() status_message.set_code(jina_pb2.StatusProto.SUCCESS) return status_message.proto except Exception as ex: status_message = StatusMessage() status_message.set_exception(ex) return status_message.proto
async def recv_dry_run(self): """Receive dry run response in bytes from server ..note:: aiohttp allows only one task which can `receive` concurrently. we need to make sure we don't create multiple tasks with `recv_message` :yield: response objects received from server """ async for response in self.response_iter: yield StatusMessage(response.data)
async def _flow_health(): """ Get the health of the complete Flow service. .. # noqa: DAR201 """ da = DocumentArray() try: _ = await _get_singleton_result( request_generator( exec_endpoint=__dry_run_endpoint__, data=da, data_type=DataInputType.DOCUMENT, )) status_message = StatusMessage() status_message.set_code(jina_pb2.StatusProto.SUCCESS) return status_message.to_dict() except Exception as ex: status_message = StatusMessage() status_message.set_exception(ex) return status_message.to_dict(use_integers_for_enums=True)
async def websocket_endpoint( websocket: WebSocket, response: Response ): # 'response' is a FastAPI response, not a Jina response from jina.proto import jina_pb2 from jina.serve.executors import __dry_run_endpoint__ await manager.connect(websocket) da = DocumentArray() try: async for _ in streamer.stream( request_iterator=request_generator( exec_endpoint=__dry_run_endpoint__, data=da, data_type=DataInputType.DOCUMENT, ) ): pass status_message = StatusMessage() status_message.set_code(jina_pb2.StatusProto.SUCCESS) await manager.send(websocket, status_message) except InternalNetworkError as err: manager.disconnect(websocket) msg = ( err.details() if _fits_ws_close_msg(err.details()) # some messages are too long else f'Network error while connecting to deployment at {err.dest_addr}. It may be down.' ) await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=msg) except WebSocketDisconnect: logger.info('Client successfully disconnected from server') manager.disconnect(websocket) except Exception as ex: manager.disconnect(websocket) status_message = StatusMessage() status_message.set_exception(ex) await manager.send(websocket, status_message)
def test_set_code(status_code): status = StatusMessage() status.set_code(status_code) assert status.proto.code == status_code
def test_init_fail(): with pytest.raises(BadRequestType): StatusMessage(status_object=5)
def test_init(status_pb): assert StatusMessage(status_object=None) assert StatusMessage(status_object=status_pb) assert StatusMessage(status_object=MessageToDict(status_pb)) assert StatusMessage(status_object=MessageToJson(status_pb))