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 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 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 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"))
