def test_init_ignore(self):
closer = utils.SocketCloser('sock', ignore=[TestException])
self.assertEqual(closer.sock, 'sock')
self.assertEqual(closer.err_thresh, 0)
self.assertEqual(closer.errors, None)
self.assertEqual(closer.ignore, frozenset([TestException]))
def test_init_err_thresh(self):
closer = utils.SocketCloser('sock', 10)
self.assertEqual(closer.sock, 'sock')
self.assertEqual(closer.err_thresh, 10)
self.assertEqual(closer.errors, 0)
self.assertEqual(closer.ignore, frozenset())
def test_together_noerror(self):
sock = mock.Mock()
with utils.SocketCloser(sock):
sock.make_call()
sock.make_call.assert_called_once_with()
self.assertFalse(sock.close.called)
def test_exit_closeerror(self):
sock = mock.Mock(**{'close.side_effect': socket.error()})
closer = utils.SocketCloser(sock)
result = closer.__exit__(TestException, TestException('foo'), [])
self.assertEqual(result, None)
sock.close.assert_called_once_with()
self.assertEqual(closer.errors, None)
def test_exit_ignore(self):
sock = mock.Mock()
closer = utils.SocketCloser(sock, ignore=[TestException])
result = closer.__exit__(TestException, TestException('foo'), [])
self.assertEqual(result, True)
self.assertFalse(sock.close.called)
self.assertEqual(closer.errors, None)
def test_exit_noerror(self):
sock = mock.Mock()
closer = utils.SocketCloser(sock)
result = closer.__exit__(None, None, None)
self.assertEqual(result, None)
self.assertFalse(sock.close.called)
self.assertEqual(closer.errors, None)
def test_together_closeerror(self):
sock = mock.Mock(**{'close.side_effect': socket.error()})
with self.assertRaises(TestException):
with utils.SocketCloser(sock):
sock.make_call()
raise TestException('spam')
sock.make_call.assert_called_once_with()
sock.close.assert_called_once_with()
def test_together_kill(self):
sock = mock.Mock()
with self.assertRaises(gevent.GreenletExit):
with utils.SocketCloser(sock):
sock.make_call()
raise gevent.GreenletExit()
sock.make_call.assert_called_once_with()
sock.close.assert_called_once_with()
def test_exit_kill_thresh(self):
sock = mock.Mock()
closer = utils.SocketCloser(sock, 10)
result = closer.__exit__(gevent.GreenletExit, gevent.GreenletExit(),
[])
self.assertEqual(result, None)
sock.close.assert_called_once_with()
self.assertEqual(closer.errors, 0)
def connect(self, target, acceptor, wrapper=None):
"""
Initiate a connection from the tendril manager's endpoint.
Once the connection is completed, a TCPTendril object will be
created and passed to the given acceptor.
:param target: The target of the connection attempt.
:param acceptor: A callable which will initialize the state of
the new TCPTendril object.
:param wrapper: A callable taking, as its first argument, a
socket.socket object. The callable must
return a valid proxy for the socket.socket
object, which will subsequently be used to
communicate on the connection.
For passing extra arguments to the acceptor or the wrapper,
see the ``TendrilPartial`` class; for chaining together
multiple wrappers, see the ``WrapperChain`` class.
"""
# Call some common sanity-checks
super(TCPTendrilManager, self).connect(target, acceptor, wrapper)
# Set up the socket
sock = socket.socket(self.addr_family, socket.SOCK_STREAM)
with utils.SocketCloser(sock, ignore=[application.RejectConnection]):
# Bind to our endpoint
sock.bind(self.endpoint)
# Connect to our target
sock.connect(target)
# Call any wrappers
if wrapper:
sock = wrapper(sock)
# Now, construct a Tendril
tend = TCPTendril(self, sock)
# Finally, set up the application
tend.application = acceptor(tend)
# OK, let's track the tendril
self._track_tendril(tend)
# Start the tendril
tend._start()
# Might as well return the tendril, too
return tend
# The acceptor raised a RejectConnection exception, apparently
sock.close()
return None
def test_together_noerror_thresh(self):
sock = mock.Mock()
closer = utils.SocketCloser(sock, 10)
closer.errors = 5
with closer:
sock.make_call()
sock.make_call.assert_called_once_with()
self.assertFalse(sock.close.called)
self.assertEqual(closer.errors, 4)
def test_together_kill_thresh(self):
sock = mock.Mock()
closer = utils.SocketCloser(sock, 10)
with self.assertRaises(gevent.GreenletExit):
with closer:
sock.make_call()
raise gevent.GreenletExit()
sock.make_call.assert_called_once_with()
sock.close.assert_called_once_with()
self.assertEqual(closer.errors, 0)
def test_exit_error_thresh(self):
sock = mock.Mock()
closer = utils.SocketCloser(sock, 1)
result = closer.__exit__(TestException, TestException('foo'), [])
self.assertEqual(result, True)
self.assertFalse(sock.close.called)
self.assertEqual(closer.errors, 1)
result = closer.__exit__(TestException, TestException('foo'), [])
self.assertEqual(result, None)
sock.close.assert_called_once_with()
self.assertEqual(closer.errors, 1)
def test_together_closeerror_thresh(self):
sock = mock.Mock(**{'close.side_effect': socket.error()})
closer = utils.SocketCloser(sock, 1)
with closer:
sock.make_call()
raise TestException('spam')
sock.make_call.assert_called_once_with()
self.assertFalse(sock.close.called)
self.assertEqual(closer.errors, 1)
sock.reset_mock()
with self.assertRaises(TestException):
with closer:
sock.make_call()
raise TestException('spam')
sock.make_call.assert_called_once_with()
sock.close.assert_called_once_with()
self.assertEqual(closer.errors, 1)
def listener(self, acceptor, wrapper):
"""
Listens for new connections to the manager's endpoint. Once a
new connection is received, a TCPTendril object is generated
for it and it is passed to the acceptor, which must initialize
the state of the connection. If no acceptor is given, no new
connections can be initialized.
:param acceptor: If given, specifies a callable that will be
called with each newly received TCPTendril;
that callable is responsible for initial
acceptance of the connection and for setting
up the initial state of the connection. If
not given, no new connections will be
accepted by the TCPTendrilManager.
:param wrapper: A callable taking, as its first argument, a
socket.socket object. The callable must
return a valid proxy for the socket.socket
object, which will subsequently be used to
communicate on the connection.
"""
# If we have no acceptor, there's nothing for us to do here
if not acceptor:
# Not listening on anything
self.local_addr = None
# Just sleep in a loop
while True:
gevent.sleep(600)
return # Pragma: nocover
# OK, set up the socket
sock = socket.socket(self.addr_family, socket.SOCK_STREAM)
with utils.SocketCloser(sock):
# Set up SO_REUSEADDR
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
# Bind to our endpoint
sock.bind(self.endpoint)
# Get the assigned port number
self.local_addr = sock.getsockname()
# Call any wrappers
if wrapper:
sock = wrapper(sock)
# Initiate listening
sock.listen(self.backlog)
# OK, now go into an accept loop with an error threshold of 10
closer = utils.SocketCloser(sock,
10,
ignore=[application.RejectConnection])
while True:
with closer:
cli, addr = sock.accept()
# OK, the connection has been accepted; construct a
# Tendril for it
tend = TCPTendril(self, cli, addr)
# Set up the application
with utils.SocketCloser(cli):
tend.application = acceptor(tend)
# Make sure we track the new tendril, but only if
# the acceptor doesn't throw any exceptions
self._track_tendril(tend)
# Start the tendril
tend._start()
def test_enter(self):
closer = utils.SocketCloser('sock')
result = closer.__enter__()
self.assertEqual(closer, result)
