Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/bots/service.py: 92.86%

1# Copyright (C) 2018 Bloomberg LP 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# <http://www.apache.org/licenses/LICENSE-2.0> 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15 

16""" 

17BotsService 

18================= 

19 

20""" 

21 

22from typing import Dict, cast 

23 

24import grpc 

25from google.protobuf import empty_pb2 

26 

27from buildgrid._protos.google.devtools.remoteworkers.v1test2.bots_pb2 import DESCRIPTOR as BOTS_DESCRIPTOR 

28from buildgrid._protos.google.devtools.remoteworkers.v1test2.bots_pb2 import ( 

29 BotSession, 

30 CreateBotSessionRequest, 

31 PostBotEventTempRequest, 

32 UpdateBotSessionRequest, 

33) 

34from buildgrid._protos.google.devtools.remoteworkers.v1test2.bots_pb2_grpc import ( 

35 BotsServicer, 

36 add_BotsServicer_to_server, 

37) 

38from buildgrid.server.bots.instance import BotsInterface 

39from buildgrid.server.decorators import rpc 

40from buildgrid.server.enums import BotStatus 

41from buildgrid.server.exceptions import InvalidArgumentError 

42from buildgrid.server.servicer import InstancedServicer 

43from buildgrid.server.utils.cancellation import CancellationContext 

44 

45 

46def _instance_name_from_bot_name(name: str) -> str: 

47 names = name.split("/") 

48 return "/".join(names[:-1]) 

49 

50 

51class BotsService(BotsServicer, InstancedServicer[BotsInterface]): 

52 SERVICE_NAME = "Bots" 

53 REGISTER_METHOD = add_BotsServicer_to_server 

54 FULL_NAME = BOTS_DESCRIPTOR.services_by_name[SERVICE_NAME].full_name 

55 

56 @rpc(instance_getter=lambda r: cast(str, r.parent)) 

57 def CreateBotSession(self, request: CreateBotSessionRequest, context: grpc.ServicerContext) -> BotSession: 

58 return self.current_instance.create_bot_session( 

59 request.bot_session, CancellationContext(context), deadline=context.time_remaining() 

60 ) 

61 

62 @rpc(instance_getter=lambda r: _instance_name_from_bot_name(r.name)) 

63 def UpdateBotSession(self, request: UpdateBotSessionRequest, context: grpc.ServicerContext) -> BotSession: 

64 if request.name != request.bot_session.name: 

65 raise InvalidArgumentError( 

66 "Name in UpdateBotSessionRequest does not match BotSession. " 

67 f" UpdateBotSessionRequest.name=[{request.name}] BotSession.name=[{request.bot_session.name}]" 

68 ) 

69 

70 bot_session, metadata = self.current_instance.update_bot_session( 

71 request.bot_session, CancellationContext(context), deadline=context.time_remaining() 

72 ) 

73 context.set_trailing_metadata(metadata) # type: ignore[arg-type] # tricky covariance issue. 

74 

75 return bot_session 

76 

77 @rpc(instance_getter=lambda r: _instance_name_from_bot_name(r.name)) 

78 def PostBotEventTemp(self, request: PostBotEventTempRequest, context: grpc.ServicerContext) -> empty_pb2.Empty: 

79 context.set_code(grpc.StatusCode.UNIMPLEMENTED) 

80 return empty_pb2.Empty() 

81 

82 def query_connected_bots_for_instance(self, instance_name: str) -> int: 

83 if instance := self.instances.get(instance_name): 

84 return instance.scheduler.job_assigner.listener_count(instance_name) 

85 return 0 

86 

87 def count_bots_by_status(self, instance_name: str) -> Dict[BotStatus, int]: 

88 if instance := self.instances.get(instance_name): 

89 return instance.count_bots_by_status() 

90 return {}