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