1# Copyright 2016 gRPC authors. 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"""The Python client used to test negative http2 conditions.""" 15 16import argparse 17 18import grpc 19import time 20from src.proto.grpc.testing import test_pb2_grpc 21from src.proto.grpc.testing import messages_pb2 22 23 24def _validate_payload_type_and_length(response, expected_type, expected_length): 25 if response.payload.type is not expected_type: 26 raise ValueError('expected payload type %s, got %s' % 27 (expected_type, type(response.payload.type))) 28 elif len(response.payload.body) != expected_length: 29 raise ValueError('expected payload body size %d, got %d' % 30 (expected_length, len(response.payload.body))) 31 32 33def _expect_status_code(call, expected_code): 34 if call.code() != expected_code: 35 raise ValueError('expected code %s, got %s' % (expected_code, 36 call.code())) 37 38 39def _expect_status_details(call, expected_details): 40 if call.details() != expected_details: 41 raise ValueError('expected message %s, got %s' % (expected_details, 42 call.details())) 43 44 45def _validate_status_code_and_details(call, expected_code, expected_details): 46 _expect_status_code(call, expected_code) 47 _expect_status_details(call, expected_details) 48 49 50# common requests 51_REQUEST_SIZE = 314159 52_RESPONSE_SIZE = 271828 53 54_SIMPLE_REQUEST = messages_pb2.SimpleRequest( 55 response_type=messages_pb2.COMPRESSABLE, 56 response_size=_RESPONSE_SIZE, 57 payload=messages_pb2.Payload(body=b'\x00' * _REQUEST_SIZE)) 58 59 60def _goaway(stub): 61 first_response = stub.UnaryCall(_SIMPLE_REQUEST) 62 _validate_payload_type_and_length(first_response, messages_pb2.COMPRESSABLE, 63 _RESPONSE_SIZE) 64 time.sleep(1) 65 second_response = stub.UnaryCall(_SIMPLE_REQUEST) 66 _validate_payload_type_and_length(second_response, 67 messages_pb2.COMPRESSABLE, _RESPONSE_SIZE) 68 69 70def _rst_after_header(stub): 71 resp_future = stub.UnaryCall.future(_SIMPLE_REQUEST) 72 _validate_status_code_and_details(resp_future, grpc.StatusCode.INTERNAL, 73 "Received RST_STREAM with error code 0") 74 75 76def _rst_during_data(stub): 77 resp_future = stub.UnaryCall.future(_SIMPLE_REQUEST) 78 _validate_status_code_and_details(resp_future, grpc.StatusCode.INTERNAL, 79 "Received RST_STREAM with error code 0") 80 81 82def _rst_after_data(stub): 83 resp_future = stub.UnaryCall.future(_SIMPLE_REQUEST) 84 _validate_status_code_and_details(resp_future, grpc.StatusCode.INTERNAL, 85 "Received RST_STREAM with error code 0") 86 87 88def _ping(stub): 89 response = stub.UnaryCall(_SIMPLE_REQUEST) 90 _validate_payload_type_and_length(response, messages_pb2.COMPRESSABLE, 91 _RESPONSE_SIZE) 92 93 94def _max_streams(stub): 95 # send one req to ensure server sets MAX_STREAMS 96 response = stub.UnaryCall(_SIMPLE_REQUEST) 97 _validate_payload_type_and_length(response, messages_pb2.COMPRESSABLE, 98 _RESPONSE_SIZE) 99 100 # give the streams a workout 101 futures = [] 102 for _ in range(15): 103 futures.append(stub.UnaryCall.future(_SIMPLE_REQUEST)) 104 for future in futures: 105 _validate_payload_type_and_length( 106 future.result(), messages_pb2.COMPRESSABLE, _RESPONSE_SIZE) 107 108 109def _run_test_case(test_case, stub): 110 if test_case == 'goaway': 111 _goaway(stub) 112 elif test_case == 'rst_after_header': 113 _rst_after_header(stub) 114 elif test_case == 'rst_during_data': 115 _rst_during_data(stub) 116 elif test_case == 'rst_after_data': 117 _rst_after_data(stub) 118 elif test_case == 'ping': 119 _ping(stub) 120 elif test_case == 'max_streams': 121 _max_streams(stub) 122 else: 123 raise ValueError("Invalid test case: %s" % test_case) 124 125 126def _args(): 127 parser = argparse.ArgumentParser() 128 parser.add_argument( 129 '--server_host', 130 help='the host to which to connect', 131 type=str, 132 default="127.0.0.1") 133 parser.add_argument( 134 '--server_port', 135 help='the port to which to connect', 136 type=int, 137 default="8080") 138 parser.add_argument( 139 '--test_case', 140 help='the test case to execute', 141 type=str, 142 default="goaway") 143 return parser.parse_args() 144 145 146def _stub(server_host, server_port): 147 target = '{}:{}'.format(server_host, server_port) 148 channel = grpc.insecure_channel(target) 149 grpc.channel_ready_future(channel).result() 150 return test_pb2_grpc.TestServiceStub(channel) 151 152 153def main(): 154 args = _args() 155 stub = _stub(args.server_host, args.server_port) 156 _run_test_case(args.test_case, stub) 157 158 159if __name__ == '__main__': 160 main() 161