def test_idle_does_not_create_response(command, response_expected):
args = set_pea_parser().parse_args([])
with Pea(args) as p:
msg = ControlMessage(command, pod_name='fake_pod')
with zmq.Context().socket(zmq.PAIR) as socket:
socket.connect(f'tcp://localhost:{p.args.port_ctrl}')
socket.send_multipart(msg.dump())
assert socket.poll(timeout=1000) == response_expected
def test_control_msg(command):
msg = ControlMessage(command)
assert msg.proto.envelope.request_type == 'ControlRequest'
assert msg.request.control.command == getattr(
jina_pb2.RequestProto.ControlRequestProto, command
)
assert msg.request.command == command
def _create_msg(args):
msg = ControlMessage('STATUS')
routing_pb = jina_pb2.RoutingTableProto()
routing_table = {
'active_pod': 'pod1',
'pods': {
'pod1': {
'host': '0.0.0.0',
'port': args.port_in,
'expected_parts': 1,
'out_edges': [{
'pod': 'pod2'
}],
},
'pod2': {
'host': '0.0.0.0',
'port': args.port_in,
'expected_parts': 1,
'out_edges': [],
},
},
}
json_format.ParseDict(routing_table, routing_pb)
msg.envelope.routing_table.CopyFrom(routing_pb)
return msg
def send_ctrl_msg(pod_address: str, command: str, timeout=1.0):
"""
Sends a control message via gRPC to pod_address
:param pod_address: the pod to send the command to
:param command: the command to send (TERMINATE/ACTIVATE/...)
:param timeout: optional timeout for the request in seconds
:returns: Empty protobuf struct
"""
stub = Grpclet._create_grpc_stub(pod_address, is_async=False)
response = stub.Call(ControlMessage(command), timeout=timeout)
return response
def test_connection_pool_same_host(mocker):
create_mock = mocker.Mock()
send_mock = mocker.Mock()
pool = ConnectionPool()
pool._create_connection = create_mock
pool._send_message = send_mock
pool.send_message(msg=ControlMessage(command='IDLE'), target_address='1.1.1.1:53')
assert send_mock.call_count == 1
assert create_mock.call_count == 1
pool.send_message(msg=ControlMessage(command='IDLE'), target_address='1.1.1.1:53')
assert send_mock.call_count == 2
assert create_mock.call_count == 1
pool.send_message(msg=ControlMessage(command='IDLE'), target_address='1.1.1.1:54')
print(send_mock.call_count, flush=True)
assert send_mock.call_count == 3
assert create_mock.call_count == 2
pool.close()
def test_bad_control_command():
with pytest.raises(ValueError):
ControlMessage('hello world')