def test_procedure_call(self):
# Create call to be streamed
call = KRPC.ProcedureCall()
call.service = 'TestService'
call.procedure = 'Counter'
# Call AddStream
add_stream = KRPC.ProcedureCall()
add_stream.service = 'KRPC'
add_stream.procedure = 'AddStream'
add_stream_arg = KRPC.Argument()
add_stream_arg.position = 0
add_stream_arg.value = call.SerializeToString()
add_stream.arguments.extend([add_stream_arg])
request = KRPC.Request()
request.calls.extend([add_stream])
self.rpc_send(request)
response = self.rpc_recv(KRPC.Response)
self.assertEqual(1, len(response.results))
stream = KRPC.Stream()
stream.ParseFromString(response.results[0].value)
# Receive some stream updates
for expected in range(1, 10):
update = self.stream_recv(KRPC.StreamUpdate)
self.assertEqual(1, len(update.results))
self.assertEqual(stream.id, update.results[0].id)
actual = self.decode_int32(update.results[0].result.value)
self.assertEqual(expected, actual)
def test_get_client_name(self):
call = KRPC.ProcedureCall()
call.service = 'KRPC'
call.procedure = 'GetClientName'
request = KRPC.MultiplexedRequest()
request.request.calls.extend([call])
self.rpc_send(request)
response = self.rpc_recv(KRPC.MultiplexedResponse).response
self.assertEqual('krpcserialio',
self.decode_string(response.results[0].value))
def test_procedure_call(self):
call = KRPC.ProcedureCall()
call.service = 'KRPC'
call.procedure = 'GetStatus'
request = KRPC.Request()
request.calls.extend([call])
self.rpc_send(request)
response = self.rpc_recv(KRPC.Response)
msg = KRPC.Status()
msg.ParseFromString(response.results[0].value)
self.assertNotEqual('', msg.version)
def test_procedure_call_with_arg(self):
call = KRPC.ProcedureCall()
call.service = 'TestService'
call.procedure = 'Int32ToString'
arg = KRPC.Argument()
arg.position = 0
arg.value = self.encode_int32(42)
call.arguments.extend([arg])
request = KRPC.Request()
request.calls.extend([call])
self.rpc_send(request)
response = self.rpc_recv(KRPC.Response)
self.assertEqual('42', self.decode_string(response.results[0].value))
def test_named_client(self):
self.connect(self.address(),
self.rpc_port(),
None,
name='TheClientName')
call = KRPC.ProcedureCall()
call.service = 'KRPC'
call.procedure = 'GetClientName'
request = KRPC.Request()
request.calls.extend([call])
self.rpc_send(request)
response = self.rpc_recv(KRPC.Response)
self.assertEqual('TheClientName',
self.decode_string(response.results[0].value))
def _invoke(self, service, procedure, args, param_names, param_types,
return_type):
""" Execute an RPC """
# Build the request
call = self._build_call(service, procedure, args, param_names,
param_types, return_type)
request = KRPC.Request()
request.calls.extend([call])
# Send the request
with self._rpc_connection_lock:
self._rpc_connection.send_message(request)
response = self._rpc_connection.receive_message(KRPC.Response)
# Check for an error response
if response.HasField('error'):
raise self._build_error(response.error)
# Check for an error in the procedure results
if response.results[0].HasField('error'):
raise self._build_error(response.results[0].error)
# Decode the response and return the (optional) result
result = None
if return_type is not None:
result = Decoder.decode(response.results[0].value, return_type)
if isinstance(result, KRPC.Event):
result = Event(self, result)
return result
def test_encode_message(self):
request = KRPC.Request()
request.service = 'ServiceName'
request.procedure = 'ProcedureName'
data = Encoder.encode(request, self.types.as_type('KRPC.Request'))
expected = '0a0b536572766963654e616d65120d50726f6365647572654e616d65'
self.assertEqual(expected, hexlify(data))
def connect(cls, port_name):
cls.rpc_conn = serial.Serial(port_name)
request = KRPC.MultiplexedRequest()
request.connection_request.type = KRPC.ConnectionRequest.RPC
request.connection_request.client_name = 'krpcserialio'
cls.send(cls.rpc_conn, request)
response = cls.recv(cls.rpc_conn, KRPC.ConnectionResponse)
assert response.status == KRPC.ConnectionResponse.OK
def _protobuf_type(code, service=None, name=None, types=None):
protobuf_type = KRPC.Type()
protobuf_type.code = code
if service is not None:
protobuf_type.service = service
if name is not None:
protobuf_type.name = name
if types is not None:
protobuf_type.types.extend(types)
return protobuf_type
def encode(cls, x, typ):
""" Encode a message or value of the given protocol buffer type """
if isinstance(typ, MessageType):
return x.SerializeToString()
elif isinstance(typ, ValueType):
return cls._encode_value(x, typ)
elif isinstance(typ, EnumerationType):
return cls._encode_value(x.value, cls._types.sint32_type)
elif isinstance(typ, ClassType):
object_id = x._object_id if x is not None else 0
return cls._encode_value(object_id, cls._types.uint64_type)
elif isinstance(typ, ListType):
msg = KRPC.List()
msg.items.extend(cls.encode(item, typ.value_type) for item in x)
return msg.SerializeToString()
elif isinstance(typ, DictionaryType):
msg = KRPC.Dictionary()
entries = []
for key, value in sorted(x.items(), key=lambda i: i[0]):
entry = KRPC.DictionaryEntry()
entry.key = cls.encode(key, typ.key_type)
entry.value = cls.encode(value, typ.value_type)
entries.append(entry)
msg.entries.extend(entries)
return msg.SerializeToString()
elif isinstance(typ, SetType):
msg = KRPC.Set()
msg.items.extend(cls.encode(item, typ.value_type) for item in x)
return msg.SerializeToString()
elif isinstance(typ, TupleType):
msg = KRPC.Tuple()
if len(x) != len(typ.value_types):
raise EncodingError('Tuple has wrong number of elements. ' +
'Expected %d, got %d.' %
(len(typ.value_types), len(x)))
msg.items.extend(
cls.encode(item, value_type)
for item, value_type in zip(x, typ.value_types))
return msg.SerializeToString()
else:
raise EncodingError('Cannot encode objects of type ' +
str(type(x)))
def connect(cls, address, rpc_port, stream_port, name=None):
if name is not None:
name = '?name=' + name
else:
name = ''
cls.rpc_conn = websocket.create_connection('ws://%s:%d/%s' %
(address, rpc_port, name))
if stream_port is None:
cls.stream_conn = None
else:
call = KRPC.ProcedureCall()
call.service = 'KRPC'
call.procedure = 'GetClientID'
request = KRPC.Request()
request.calls.extend([call])
cls.send(cls.rpc_conn, request)
response = cls.recv(cls.rpc_conn, KRPC.Response)
client_identifier = cls.decode_bytes(response.results[0].value)
client_identifier = base64.b64encode(client_identifier)
cls.stream_conn = websocket.create_connection(
'ws://%s:%d/?id=%s' %
(address, stream_port, client_identifier))
def _build_request(self, service, procedure, args,
param_names, param_types, return_type):
""" Build a KRPC.Request object """
request = KRPC.Request(service=service, procedure=procedure)
for i, (value, typ) in enumerate(itertools.izip(args, param_types)):
if isinstance(value, DefaultArgument):
continue
if not isinstance(value, typ.python_type):
try:
value = self._types.coerce_to(value, typ)
except ValueError:
raise TypeError(
'%s.%s() argument %d must be a %s, got a %s' %
(service, procedure, i, typ.python_type, type(value)))
request.arguments.add(position=i, value=Encoder.encode(value, typ))
return request
def _build_call(self, service, procedure, args, param_names, param_types,
return_type):
# pylint: disable=unused-argument
""" Build a KRPC.ProcedureCall object """
call = KRPC.ProcedureCall()
call.service = service
call.procedure = procedure
for i, (value, typ) in enumerate(itertools.izip(args, param_types)):
if isinstance(value, DefaultArgument):
continue
if not isinstance(value, typ.python_type):
try:
value = self._types.coerce_to(value, typ)
except ValueError:
raise TypeError(
'%s.%s() argument %d must be a %s, got a %s' %
(service, procedure, i, typ.python_type, type(value)))
call.arguments.add(position=i, value=Encoder.encode(value, typ))
return call
except IndexError:
pass
# Receive the message data
data = conn.recv(size)
# Decode the message
msg = msg_type()
msg.ParseFromString(data)
return msg
# Open a TCP/IP socket to the RPC server
rpc_conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
rpc_conn.connect(('127.0.0.1', 50000))
# Send an RPC connection request
request = KRPC.ConnectionRequest()
request.type = KRPC.ConnectionRequest.RPC
request.client_name = 'Jeb'
send_message(rpc_conn, request)
# Receive the connection response
response = recv_message(rpc_conn, KRPC.ConnectionResponse)
# Check the connection was successful
if response.status != KRPC.ConnectionResponse.OK:
raise RuntimeError('Connection failed: ' + response.message)
print('Connected to RPC server')
# Invoke the KRPC.GetStatus RPC
call = KRPC.ProcedureCall()
call.service = 'KRPC'
except IndexError:
pass
# Receive the message data
data = conn.recv(size)
# Decode the message
msg = msg_type()
msg.ParseFromString(data)
return msg
# Open a TCP/IP socket to the RPC server
rpc_conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
rpc_conn.connect(('127.0.0.1', 50000))
# Send an RPC connection request
request = KRPC.ConnectionRequest()
request.type = KRPC.ConnectionRequest.RPC
request.client_name = 'Jeb'
send_message(rpc_conn, request)
# Receive the connection response
response = recv_message(rpc_conn, KRPC.ConnectionResponse)
# Check the connection was successful
if response.status != KRPC.ConnectionResponse.OK:
raise RuntimeError('Connection failed: ' + response.message)
print('Connected to RPC server')
# Print out the clients identifier
print('RPC client idenfitier = %s' %
binascii.hexlify(bytearray(response.client_identifier)).decode('utf8'))