def test_retry_on_error_success():
mock_should_retry_on_error = mock.patch('tchannel.tornado.Request.should_retry_on_error') # noqa
mock_should_retry_on_error.return_value = True
endpoint = 'tchannelretrytest'
tchannel = yield chain(2, endpoint)
hook = MyTestHook()
tchannel.hooks.register(hook)
tchannel_success = TChannel(name='test', hostport='localhost:0')
tchannel_success.register(endpoint, 'raw', handler_success)
tchannel_success.listen()
tchannel.peers.get(tchannel_success.hostport)
response = yield tchannel.request().send(
endpoint,
"test",
"test",
headers={
're': retry.CONNECTION_ERROR_AND_TIMEOUT,
},
ttl=1,
retry_limit=2,
)
header = yield response.get_header()
body = yield response.get_body()
assert body == b"success"
assert header == b""
assert hook.received_response == 1
assert hook.received_error == 2
def test_retry_on_error_success():
endpoint = b'tchannelretrytest'
tchannel = chain(2, endpoint)
tchannel_success = TChannel(name='test', hostport='localhost:0')
tchannel_success.register(endpoint, 'raw', handler_success)
tchannel_success.listen()
tchannel.peers.get(tchannel_success.hostport)
with (
patch(
'tchannel.tornado.Request.should_retry_on_error',
autospec=True)
) as mock_should_retry_on_error:
mock_should_retry_on_error.return_value = True
response = yield tchannel.request(
score_threshold=0
).send(
endpoint,
"test",
"test",
headers={
're': RetryType.CONNECTION_ERROR_AND_TIMEOUT,
},
ttl=0.01,
attempt_times=3,
retry_delay=0.01,
)
header = yield response.get_header()
body = yield response.get_body()
assert body == "success"
assert header == ""
def test_retry_on_error_success(mock_should_retry_on_error):
mock_should_retry_on_error.return_value = True
endpoint = b'tchannelretrytest'
tchannel = yield chain(2, endpoint)
hook = MyTestHook()
tchannel.hooks.register(hook)
tchannel_success = TChannel(name='test', hostport='localhost:0')
tchannel_success.register(endpoint, 'raw', handler_success)
tchannel_success.listen()
tchannel.peers.get(tchannel_success.hostport)
response = yield tchannel.request().send(
endpoint,
"test",
"test",
headers={
're': retry.CONNECTION_ERROR_AND_TIMEOUT,
},
ttl=1,
retry_limit=2,
)
header = yield response.get_header()
body = yield response.get_body()
assert body == "success"
assert header == ""
assert hook.received_response == 1
assert hook.received_error == 2
def main(): # pragma: no cover
import sys
parser = argparse.ArgumentParser()
parser.add_argument(
"--port",
dest="port", default=8888, type=int,
)
parser.add_argument(
"--host",
dest="host", default="127.0.0.1"
)
args = parser.parse_args()
client = TChannel('%s:%d' % (args.host, args.port))
# TODO: service=test_as_raw
handler = RequestDispatcher()
client.host(handler)
# TODO: do we need to implement these differently?
handler.register("echo", echo)
handler.register("streaming_echo", echo)
client.listen()
print 'listening on', client.hostport
sys.stdout.flush()
tornado.ioloop.IOLoop.instance().start()
def main(): # pragma: no cover
args = get_args()
client = TChannel('localhost:%d' % args.port)
register_example_endpoints(client)
client.listen()
tornado.ioloop.IOLoop.instance().start()
def main(): # pragma: no cover
args = get_args()
client = TChannel(
name='tchannel_server',
hostport='%s:%d' % (args.host, args.port),
)
register_example_endpoints(client)
client.listen()
tornado.ioloop.IOLoop.instance().start()
def main():
args = get_args()
app = TChannel(
name='raw-server',
hostport='%s:%d' % (args.host, args.port),
)
register_example_endpoints(app)
app.listen()
print("listening on %s" % app.hostport)
sys.stdout.flush()
tornado.ioloop.IOLoop.instance().start()
def test_advertise():
server = TChannel(name="test_server")
@server.register('ad', 'json')
@tornado.gen.coroutine
def ad(request, response):
body = yield request.get_body()
response.write_body(body)
server.listen()
channel = TChannel(name='test')
response = yield hyperbahn.advertise(channel, 'test', [server.hostport])
result = yield response.get_body()
assert (result == '{"services": [{"serviceName": "test", "cost": 0}]}'
or result == '{"services": [{"cost": 0, "serviceName": "test"}]}')
def main():
args = get_args()
app = TChannel(
name='json-server',
hostport='%s:%d' % (args.host, args.port),
)
register_example_endpoints(app)
def say_hi_json(request, response, proxy):
response.write_body({'hi': 'Hello, world!'})
app.register(endpoint="hi-json", scheme="json", handler=say_hi_json)
app.listen()
tornado.ioloop.IOLoop.instance().start()
def test_advertise():
server = TChannel(name="test_server")
@server.register("ad", "json")
@tornado.gen.coroutine
def ad(request, response):
body = yield request.get_body()
response.write_body(body)
server.listen()
channel = TChannel(name="test")
response = yield hyperbahn.advertise(channel, "test", [server.hostport])
result = yield response.get_body()
assert (
result == '{"services": [{"serviceName": "test", "cost": 0}]}'
or result == '{"services": [{"cost": 0, "serviceName": "test"}]}'
)
def test_false_result(thrift_service):
# Verify that we aren't treating False as None.
app = TChannel(name='app')
@app.register(thrift_service)
def healthy(request, response):
return False
app.listen()
client = TChannel(name='client')
response = yield client.request(hostport=app.hostport,
arg_scheme='thrift').send(
'Service::healthy', '\x00\x00', '\x00')
body = yield response.get_body()
assert body == '\x02\x00\x00\x00\x00'
def test_false_result(thrift_service):
# Verify that we aren't treating False as None.
app = TChannel(name='app')
@app.register(thrift_service)
def healthy(request, response):
return False
app.listen()
client = TChannel(name='client')
response = yield client.request(
hostport=app.hostport, arg_scheme='thrift'
).send('Service::healthy', '\x00\x00', '\x00')
body = yield response.get_body()
assert body == '\x02\x00\x00\x00\x00'
def test_close_stops_listening():
server = TChannel(name='server')
server.listen()
host = server.host
port = server.port
# Can connect
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.connect((host, port))
sock.close()
server.close()
# Can't connect
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
with pytest.raises(socket.error):
sock.connect((host, port))
def test_close_stops_listening():
server = TChannel(name='server')
server.listen()
host = server.host
port = server.port
# Can connect
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.connect((host, port))
sock.close()
server.close()
# Can't connect
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
with pytest.raises(socket.error):
sock.connect((host, port))
class Fakebahn(object):
def __init__(self, handle):
self.tchannel = TChannel(name='hyperbahn')
self.count = 0 # number of ad requests received
@self.tchannel.register('ad', 'json')
@gen.coroutine
def ad(request, response):
self.count += 1
yield handle(request, response)
@property
def hostport(self):
return self.tchannel.hostport
def start(self):
self.tchannel.listen()
def stop(self):
self.tchannel.close()
class Fakebahn(object):
def __init__(self, handle):
self.tchannel = TChannel(name='hyperbahn')
self.count = 0 # number of ad requests received
@self.tchannel.register('ad', 'json')
@gen.coroutine
def ad(request, response):
self.count += 1
yield handle(request, response)
@property
def hostport(self):
return self.tchannel.hostport
def start(self):
self.tchannel.listen()
def stop(self):
self.tchannel.close()
def server(endpoint):
tchannel_server = TChannel(name='testserver', hostport='localhost:0')
tchannel_server.register(endpoint, 'raw', handler_error)
tchannel_server.listen()
return tchannel_server
class VCRProxyService(object):
def __init__(self, cassette, unpatch):
"""
:param unpatch:
A function returning a context manager which temporarily unpatches
any monkey patched code so that a real request can be made.
:param cassette:
Cassette being played.
"""
self.unpatch = unpatch
self.cassette = cassette
self.tchannel = TChannel("proxy-server")
self.tchannel.register(VCRProxy, handler=self.send)
@wrap_uncaught(
reraise=(VCRProxy.CannotRecordInteractionsError, VCRProxy.RemoteServiceError, VCRProxy.VCRServiceError)
)
@gen.coroutine
def send(self, request, response):
cassette = self.cassette
request = request.args.request
# TODO decode requests and responses based on arg scheme into more
# readable formats.
# Because Thrift doesn't handle UTF-8 correctly right now
request.serviceName = request.serviceName.decode("utf-8")
request.endpoint = request.endpoint.decode("utf-8")
# TODO do we care about hostport being the same?
if cassette.can_replay(request):
vcr_response = cassette.replay(request)
raise gen.Return(vcr_response)
if cassette.write_protected:
raise VCRProxy.CannotRecordInteractionsError(
"Could not find a matching response for request %s and the "
"record mode %s prevents new interactions from being "
"recorded. Your test may be performing an uenxpected "
"request." % (str(request), cassette.record_mode)
)
arg_scheme = VCRProxy.ArgScheme.to_name(request.argScheme).lower()
with self.unpatch():
# TODO propagate other request and response parameters
# TODO might make sense to tag all VCR requests with a protocol
# header of some kind
response_future = self.tchannel.request(
service=request.serviceName, arg_scheme=arg_scheme, hostport=request.hostPort
).send(
request.endpoint,
request.headers,
request.body,
headers={h.key: h.value for h in request.transportHeaders},
)
# Don't actually yield while everything is unpatched.
try:
response = yield response_future
except TChannelError as e:
raise VCRProxy.RemoteServiceError(code=e.code, message=e.message)
response_headers = yield response.get_header()
response_body = yield response.get_body()
vcr_response = VCRProxy.Response(response.status_code, response_headers, response_body)
cassette.record(request, vcr_response)
raise gen.Return(vcr_response)
@property
def hostport(self):
return self.tchannel.hostport
def start(self):
self.tchannel.listen()
def stop(self):
self.tchannel.close()
def __enter__(self):
self.start()
return self
def __exit__(self, *args):
self.stop()
def server(endpoint):
tchannel_server = TChannel(name='testserver', hostport='localhost:0')
tchannel_server.register(endpoint, 'raw', handler_error)
tchannel_server.listen()
return tchannel_server
class VCRProxyService(object):
def __init__(self, cassette, unpatch):
"""
:param unpatch:
A function returning a context manager which temporarily unpatches
any monkey patched code so that a real request can be made.
:param cassette:
Cassette being played.
"""
self.unpatch = unpatch
self.cassette = cassette
self.io_loop = None
self.thread = None
self.tchannel = None
self._running = threading.Event()
@wrap_uncaught(reraise=(
VCRProxy.CannotRecordInteractionsError,
VCRProxy.RemoteServiceError,
VCRProxy.VCRServiceError,
))
@gen.coroutine
def send(self, request, response):
cassette = self.cassette
request = request.args.request
# TODO decode requests and responses based on arg scheme into more
# readable formats.
# Because Thrift doesn't handle UTF-8 correctly right now
request.serviceName = request.serviceName.decode('utf-8')
request.endpoint = request.endpoint.decode('utf-8')
# TODO do we care about hostport being the same?
if cassette.can_replay(request):
vcr_response = cassette.replay(request)
raise gen.Return(vcr_response)
if cassette.write_protected:
raise VCRProxy.CannotRecordInteractionsError(
'Could not find a matching response for request %s and the '
'record mode %s prevents new interactions from being '
'recorded. Your test may be performing an uenxpected '
'request.' % (str(request), cassette.record_mode))
arg_scheme = VCRProxy.ArgScheme.to_name(request.argScheme).lower()
with self.unpatch():
# TODO propagate other request and response parameters
# TODO might make sense to tag all VCR requests with a protocol
# header of some kind
response_future = self.tchannel.request(
service=request.serviceName,
arg_scheme=arg_scheme,
hostport=request.hostPort,
).send(
request.endpoint,
request.headers,
request.body,
headers={h.key: h.value
for h in request.transportHeaders},
)
# Don't actually yield while everything is unpatched.
try:
response = yield response_future
except TChannelError as e:
raise VCRProxy.RemoteServiceError(
code=e.code,
message=e.message,
)
response_headers = yield response.get_header()
response_body = yield response.get_body()
vcr_response = VCRProxy.Response(
response.status_code,
response_headers,
response_body,
)
cassette.record(request, vcr_response)
raise gen.Return(vcr_response)
@property
def hostport(self):
return self.tchannel.hostport
def _run(self):
self.io_loop = IOLoop()
self.io_loop.make_current()
self.tchannel = TChannel('proxy-server')
self.tchannel.register(VCRProxy, handler=self.send)
self.tchannel.listen()
self._running.set()
self.io_loop.start()
def start(self):
self.thread = threading.Thread(target=self._run)
self.thread.start()
self._running.wait(1)
def stop(self):
self.tchannel.close()
self.tchannel = None
self.io_loop.stop()
self.io_loop = None
self.thread.join(1) # seconds
self.thread = None
def __enter__(self):
self.start()
return self
def __exit__(self, *args):
self.stop()
