def CreateBotSession(self, request, context): """Handles CreateBotSessionRequest messages. Args: request (CreateBotSessionRequest): The incoming RPC request. context (grpc.ServicerContext): Context for the RPC call. """ self.__logger.debug("CreateBotSession request from [%s]", context.peer()) instance_name = request.parent bot_status = BotStatus(request.bot_session.status) bot_id = request.bot_session.bot_id try: instance = self._get_instance(instance_name) bot_session = instance.create_bot_session(instance_name, request.bot_session) now = timestamp_pb2.Timestamp() now.GetCurrentTime() if self._is_instrumented: self.__bots[bot_id] = now self.__bots_by_instance[instance_name].add(bot_id) if bot_status in self.__bots_by_status: self.__bots_by_status[bot_status].add(bot_id) return bot_session except InvalidArgumentError as e: self.__logger.error(e) context.set_details(str(e)) context.set_code(grpc.StatusCode.INVALID_ARGUMENT) return bots_pb2.BotSession()
def get_pb2(self): return bots_pb2.BotSession( worker=self._hardware_interface.get_worker_pb2(), status=self._status, leases=self._tenant_manager.get_leases(), bot_id=self.__bot_id, name=self.__name)
def test_create_bot_session_bot_id_fail(context, instance): bot = bots_pb2.BotSession() request = bots_pb2.CreateBotSessionRequest(parent='', bot_session=bot) instance.CreateBotSession(request, context) context.set_code.assert_called_once_with(grpc.StatusCode.INVALID_ARGUMENT)
def UpdateBotSession(self, request, context): """Handles UpdateBotSessionRequest messages. Args: request (UpdateBotSessionRequest): The incoming RPC request. context (grpc.ServicerContext): Context for the RPC call. """ self.__logger.debug("UpdateBotSession request from [%s]", context.peer()) names = request.name.split("/") bot_status = BotStatus(request.bot_session.status) bot_id = request.bot_session.bot_id try: instance_name = '/'.join(names[:-1]) instance = self._get_instance(instance_name) bot_session = instance.update_bot_session( request.name, request.bot_session, deadline=context.time_remaining()) if self._is_instrumented: self.__bots[bot_id].GetCurrentTime() if bot_id not in self.__bots_by_status[bot_status]: if bot_status == BotStatus.OK: self.__bots_by_status[BotStatus.OK].add(bot_id) self.__bots_by_status[BotStatus.UNHEALTHY].discard( bot_id) elif bot_status == BotStatus.UNHEALTHY: self.__bots_by_status[BotStatus.OK].discard(bot_id) self.__bots_by_status[BotStatus.UNHEALTHY].add(bot_id) else: self.__bots_by_instance[instance_name].remove(bot_id) del self.__bots[bot_id] return bot_session except InvalidArgumentError as e: self.__logger.error(e) context.set_details(str(e)) context.set_code(grpc.StatusCode.INVALID_ARGUMENT) except OutOfSyncError as e: self.__logger.error(e) context.set_details(str(e)) context.set_code(grpc.StatusCode.DATA_LOSS) except NotImplementedError as e: self.__logger.error(e) context.set_details(str(e)) context.set_code(grpc.StatusCode.UNIMPLEMENTED) return bots_pb2.BotSession()
def __create_bot_session(queue, remote, parent, string_bot_session): bot_session = bots_pb2.BotSession() bot_session.ParseFromString(string_bot_session) interface = BotInterface(grpc.insecure_channel(remote), TIMEOUT, TIMEOUT) result = interface.create_bot_session(parent, bot_session) queue.put(result.SerializeToString())
def __update_bot_session(queue, remote, string_bot_session, update_mask): bot_session = bots_pb2.BotSession() bot_session.ParseFromString(string_bot_session) interface = BotInterface(grpc.insecure_channel(remote), TIMEOUT, TIMEOUT) result = interface.update_bot_session(bot_session, update_mask) queue.put(result.SerializeToString())
def create_bot_session(self, parent, bot_session): def __create_bot_session(queue, remote, parent, string_bot_session): bot_session = bots_pb2.BotSession() bot_session.ParseFromString(string_bot_session) interface = BotInterface(grpc.insecure_channel(remote), TIMEOUT, TIMEOUT) result = interface.create_bot_session(parent, bot_session) queue.put(result.SerializeToString()) string_bot_session = bot_session.SerializeToString() result = run_in_subprocess(__create_bot_session, self.__remote, parent, string_bot_session) bot_session = bots_pb2.BotSession() bot_session.ParseFromString(result) return bot_session
def update_bot_session(self, bot_session, update_mask=None): def __update_bot_session(queue, remote, string_bot_session, update_mask): bot_session = bots_pb2.BotSession() bot_session.ParseFromString(string_bot_session) interface = BotInterface(grpc.insecure_channel(remote), TIMEOUT, TIMEOUT) result = interface.update_bot_session(bot_session, update_mask) queue.put(result.SerializeToString()) string_bot_session = bot_session.SerializeToString() result = run_in_subprocess(__update_bot_session, self.__remote, string_bot_session, update_mask) bot_session = bots_pb2.BotSession() bot_session.ParseFromString(result) return bot_session
def __test_create_server(queue, remote): # Open a channel to the remote server: channel = grpc.insecure_channel(remote) try: stub = remote_execution_pb2_grpc.ExecutionStub(channel) request = remote_execution_pb2.ExecuteRequest(instance_name='main') response = next(stub.Execute(request)) assert response.DESCRIPTOR is operations_pb2.Operation.DESCRIPTOR except grpc.RpcError as e: if e.code() == grpc.StatusCode.UNIMPLEMENTED: queue.put(False) except AssertionError: queue.put(False) try: stub = remote_execution_pb2_grpc.ActionCacheStub(channel) request = remote_execution_pb2.GetActionResultRequest(instance_name='main') response = stub.GetActionResult(request) assert response.DESCRIPTOR is remote_execution_pb2.ActionResult.DESCRIPTOR except grpc.RpcError as e: if e.code() == grpc.StatusCode.UNIMPLEMENTED: queue.put(False) except AssertionError: queue.put(False) try: stub = remote_execution_pb2_grpc.ContentAddressableStorageStub(channel) request = remote_execution_pb2.BatchUpdateBlobsRequest(instance_name='main') response = stub.BatchUpdateBlobs(request) assert response.DESCRIPTOR is remote_execution_pb2.BatchUpdateBlobsResponse.DESCRIPTOR except grpc.RpcError as e: if e.code() == grpc.StatusCode.UNIMPLEMENTED: queue.put(False) except AssertionError: queue.put(False) try: stub = buildstream_pb2_grpc.ReferenceStorageStub(channel) request = buildstream_pb2.GetReferenceRequest(instance_name='main') response = stub.GetReference(request) assert response.DESCRIPTOR is buildstream_pb2.GetReferenceResponse.DESCRIPTOR except grpc.RpcError as e: if e.code() == grpc.StatusCode.UNIMPLEMENTED: queue.put(False) except AssertionError: queue.put(False) try: stub = bytestream_pb2_grpc.ByteStreamStub(channel) request = bytestream_pb2.ReadRequest() response = stub.Read(request) assert next(response).DESCRIPTOR is bytestream_pb2.ReadResponse.DESCRIPTOR except grpc.RpcError as e: if e.code() == grpc.StatusCode.UNIMPLEMENTED: queue.put(False) except AssertionError: queue.put(False) try: stub = operations_pb2_grpc.OperationsStub(channel) request = operations_pb2.ListOperationsRequest(name='main') response = stub.ListOperations(request) assert response.DESCRIPTOR is operations_pb2.ListOperationsResponse.DESCRIPTOR except grpc.RpcError as e: if e.code() == grpc.StatusCode.UNIMPLEMENTED: queue.put(False) except AssertionError: queue.put(False) try: session = bots_pb2.BotSession(leases=[], bot_id="test-bot", name="test-bot", status=BotStatus.OK.value) stub = bots_pb2_grpc.BotsStub(channel) request = bots_pb2.CreateBotSessionRequest(parent='main', bot_session=session) response = stub.CreateBotSession(request) assert response.DESCRIPTOR is bots_pb2.BotSession.DESCRIPTOR except grpc.RpcError as e: if e.code() == grpc.StatusCode.UNIMPLEMENTED: queue.put(False) except AssertionError: queue.put(False) # Sleep for 5 seconds to allow a useful number of metrics to get reported time.sleep(5) queue.put(True)
def bot_session(): bot = bots_pb2.BotSession() bot.bot_id = 'ana' bot.status = BotStatus.OK.value yield bot