def test_zipkin_trace(trace_server): endpoint = b'endpoint1' zipkin_tracer = ZipkinTraceHook(dst=trace_buf) tchannel = TChannel(name='test') tchannel.hooks.register(zipkin_tracer) hostport = 'localhost:%d' % trace_server.port response = yield tchannel.request(hostport).send(InMemStream(endpoint), InMemStream(hostport), InMemStream(), traceflag=True) header = yield response.get_header() body = yield response.get_body() assert header == "from handler1" assert body == "from handler2" traces = [] for trace in trace_buf.getvalue().split("\n"): if trace: traces.append(json.loads(trace)) trace_id = traces[0][0][u'trace_id'] for trace in traces: assert trace_id == trace[0][u'trace_id'] if trace[0][u'name'] == u'endpoint2': parent_span_id = trace[0][u'parent_span_id'] else: span_id = trace[0][u'span_id'] assert parent_span_id == span_id
def test_endpoint_not_found(mock_server): endpoint = b'tchanneltest' mock_server.expect_call(endpoint).and_write(headers=endpoint, body='world') tchannel = TChannel(name='test') with pytest.raises(TChannelError): yield tchannel.request(mock_server.hostport).send( InMemStream(), InMemStream(), InMemStream())
def closed_stream(body): """Builds an in-memory stream whose entire request body is the given string. :param body: Request body for the returned Stream """ stream = InMemStream(body) stream.close() return stream
def loop1(n): results = yield [ server1.request(hostport2).send(InMemStream('hello'), InMemStream(), InMemStream()) for i in xrange(n) ] for resp in results: body = yield resp.get_body() assert body == 'hello to you too'
def loop2(n): results = yield [ server2.request(hostport1).send(InMemStream('reverse'), InMemStream(), InMemStream('foo')) for i in xrange(n) ] for resp in results: body = yield resp.get_body() assert body == 'bar'
def test_build_handler(): def call(treq, tres, tchan): assert treq.transport.headers == { 'as': 'thrift', 'cn': 'test_caller' } tres.write_header('foo', 'baar') return "world" response_header = InMemStream() response_body = InMemStream() req = Request( argstreams=[ InMemStream('hello'), InMemStream('\00\00'), # no headers InMemStream('\00'), # empty struct ], scheme=ThriftArgScheme(FakeResult), headers={'cn': 'test_caller', 'as': 'thrift'}, ) req.close_argstreams() res = Response( argstreams=[ InMemStream(), response_header, response_body, ], scheme=ThriftArgScheme(FakeResult), ) tchannel = mock.Mock() handler = build_handler(FakeResult, call) yield handler(req, res, tchannel) serialized_headers = yield response_header.read() assert serialized_headers == bytearray( [ 0x00, 0x01, # num headers = 1 0x00, 0x03, # strlen('foo') = 3 ] + list('foo') + [ 0x00, 0x04, # strlen('baar') = 4 ] + list('baar') ) serialized_body = yield response_body.read() assert serialized_body == bytearray([ 0x0b, # field type = TType.STRING 0x00, 0x00, # field ID = 0 0x00, 0x00, 0x00, 0x05, # string length = 5 ] + list("world") + [ 0x00, # end struct ]) assert 0 == res.status_code
def test_deprecated_build_handler(): def call(treq, tres): assert treq.transport.headers == { 'as': 'thrift', 'cn': 'test_caller' } tres.write_header('foo', 'baar') return "world" response_header = InMemStream() response_body = InMemStream() req = Request( argstreams=[ InMemStream('hello'), InMemStream('\00\00'), # no headers InMemStream('\00'), # empty struct ], serializer=ThriftSerializer(FakeResult), headers={'cn': 'test_caller', 'as': 'thrift'}, ) req.close_argstreams() res = Response( argstreams=[ InMemStream(), response_header, response_body, ], serializer=ThriftSerializer(FakeResult), ) handler = deprecated_build_handler(FakeResult, call) yield handler(req, res) serialized_headers = yield response_header.read() assert serialized_headers == bytearray( [ 0x00, 0x01, # num headers = 1 0x00, 0x03, # strlen('foo') = 3 ] + list('foo') + [ 0x00, 0x04, # strlen('baar') = 4 ] + list('baar') ) serialized_body = yield response_body.read() assert serialized_body == bytearray([ 0x0b, # field type = TType.STRING 0x00, 0x00, # field ID = 0 0x00, 0x00, 0x00, 0x05, # string length = 5 ] + list("world") + [ 0x00, # end struct ]) assert 0 == res.status_code
def test_tchannel_call_request_fragment(mock_server, arg2, arg3): endpoint = b'tchannelpeertest' mock_server.expect_call(endpoint).and_write(headers=endpoint, body=arg3) tchannel = TChannel(name='test') response = yield tchannel.request(mock_server.hostport).send( InMemStream(endpoint), InMemStream(arg2), InMemStream(arg3)) header = yield response.get_header() body = yield response.get_body() assert header == endpoint assert body == arg3 assert response.headers['as'] == 'raw'
def test_deprecated_build_handler_exception(): def call(treq, tres): raise FakeException('fail') response_body = mock.Mock(spec=InMemStream) req = Request( argstreams=[ InMemStream('hello'), InMemStream('\00\00'), # no headers InMemStream('\00'), # empty struct ], serializer=ThriftSerializer(FakeResult), ) req.close_argstreams() res = Response( argstreams=[ InMemStream(), InMemStream(), response_body, ], serializer=ThriftSerializer(FakeResult), ) handler = deprecated_build_handler(FakeResult, call) yield handler(req, res) response_body.write.assert_called_once_with( bytearray([ 0x0c, # field type = TType.STRUCT 0x00, 0x01, # field ID = 1 0x0b, # field type = TType.STRING 0x00, 0x01, # field ID = 1 0x00, 0x00, 0x00, 0x04, # string length = 5 ] + list(b"fail") + [ 0x00, # end exception struct 0x00, # end response struct ])) assert 1 == res.status_code
def test_build_handler_application_exception(): def call(req): raise FakeException('fail') req = Request( argstreams=[ InMemStream('hello'), InMemStream('\00\00'), # no headers InMemStream('\00'), # empty struct ], serializer=ThriftSerializer(FakeResult), ) req.close_argstreams() handler = build_handler(FakeResult, call) res = yield handler(req) assert res.status == 1
def test_InMemStream(): stream = InMemStream() yield stream.write("1") yield stream.write("2") buf = yield stream.read() assert buf == "12" yield stream.write("3") buf = yield stream.read() assert buf == "3" # check internal stream buffer. assert len(stream._stream) == 0 stream.close() with pytest.raises(StreamingError): yield stream.write("4")
def test_error_during_stream(io_loop): stream = InMemStream() try: 1 / 0 except Exception as e: stream.set_exception(e) with pytest.raises(ZeroDivisionError): yield stream.read() with pytest.raises(ZeroDivisionError): yield stream.write("a")
def stream(s): s = InMemStream(s) s.close() return s
def closed_stream(body): stream = InMemStream(body) stream.close() return stream
def handler_success(request, response): response.set_body_s(InMemStream("success"))
def handler2(request, response): response.set_body_s(InMemStream("from handler2"))