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