Merge branch 'trunk' into 9460-incremental-codingstandard

twm · web-flow · commit 67aa1604e200 · 2018-06-02T15:34:59.000-07:00
diff --git a/src/twisted/application/internet.py b/src/twisted/application/internet.py
@@ -48,7 +48,7 @@
 from twisted.internet import task
 from twisted.python.failure import Failure
 from twisted.internet.defer import (
-    CancelledError, Deferred, succeed, fail
+    CancelledError, Deferred, succeed, fail, maybeDeferred
 )
 
 from automat import MethodicalMachine
@@ -548,7 +548,8 @@ class _ClientMachine(object):
 
     _machine = MethodicalMachine()
 
-    def __init__(self, endpoint, factory, retryPolicy, clock, log):
+    def __init__(self, endpoint, factory, retryPolicy, clock,
+                 prepareConnection, log):
         """
         @see: L{ClientService.__init__}
 
@@ -566,6 +567,7 @@ def __init__(self, endpoint, factory, retryPolicy, clock, log):
         self._factory = factory
         self._timeoutForAttempt = retryPolicy
         self._clock = clock
+        self._prepareConnection = prepareConnection
         self._connectionInProgress = succeed(None)
 
         self._awaitingConnected = []
@@ -633,10 +635,35 @@ def _connect(self):
 
         self._connectionInProgress = (
             self._endpoint.connect(factoryProxy)
+            .addCallback(self._runPrepareConnection)
             .addCallback(self._connectionMade)
             .addErrback(self._connectionFailed))
 
 
+    def _runPrepareConnection(self, protocol):
+        """
+        Run any C{prepareConnection} callback with the connected protocol,
+        ignoring its return value but propagating any failure.
+
+        @param protocol: The protocol of the connection.
+        @type protocol: L{IProtocol}
+
+        @return: Either:
+
+            - A L{Deferred} that succeeds with the protocol when the
+              C{prepareConnection} callback has executed successfully.
+
+            - A L{Deferred} that fails when the C{prepareConnection} callback
+              throws or returns a failed L{Deferred}.
+
+            - The protocol, when no C{prepareConnection} callback is defined.
+        """
+        if self._prepareConnection:
+            return (maybeDeferred(self._prepareConnection, protocol)
+                    .addCallback(lambda _: protocol))
+        return protocol
+
+
     @_machine.output()
     def _resetFailedAttempts(self):
         """
@@ -1009,7 +1036,9 @@ class ClientService(service.Service, object):
     """
 
     _log = Logger()
-    def __init__(self, endpoint, factory, retryPolicy=None, clock=None):
+
+    def __init__(self, endpoint, factory, retryPolicy=None, clock=None,
+                 prepareConnection=None):
         """
         @param endpoint: A L{stream client endpoint
             <interfaces.IStreamClientEndpoint>} provider which will be used to
@@ -1029,13 +1058,36 @@ def __init__(self, endpoint, factory, retryPolicy=None, clock=None):
             this attribute will not be serialized, and the default value (the
             reactor) will be restored when deserialized.
         @type clock: L{IReactorTime}
+
+        @param prepareConnection: A single argument L{callable} that may return
+            a L{Deferred}. It will be called once with the L{protocol
+            <interfaces.IProtocol>} each time a new connection is made.  It may
+            call methods on the protocol to prepare it for use (e.g.
+            authenticate) or validate it (check its health).
+
+            The C{prepareConnection} callable may raise an exception or return
+            a L{Deferred} which fails to reject the connection.  A rejected
+            connection is not used to fire an L{Deferred} returned by
+            L{whenConnected}.  Instead, L{ClientService} handles the failure
+            and continues as if the connection attempt were a failure
+            (incrementing the counter passed to C{retryPolicy}).
+
+            L{Deferred}s returned by L{whenConnected} will not fire until
+            any L{Deferred} returned by the C{prepareConnection} callable
+            fire. Otherwise its successful return value is consumed, but
+            ignored.
+
+            Present Since Twisted NEXT
+
+        @type prepareConnection: L{callable}
+
         """
         clock = _maybeGlobalReactor(clock)
         retryPolicy = _defaultPolicy if retryPolicy is None else retryPolicy
 
         self._machine = _ClientMachine(
             endpoint, factory, retryPolicy, clock,
-            log=self._log,
+            prepareConnection=prepareConnection, log=self._log,
         )
 
 
diff --git a/src/twisted/application/test/test_internet.py b/src/twisted/application/test/test_internet.py
@@ -1031,3 +1031,148 @@ def test_stopServiceOnStoppedService(self):
 
         self.assertIsNone(self.successResultOf(firstStopDeferred))
         self.assertIsNone(self.successResultOf(secondStopDeferred))
+
+
+    def test_prepareConnectionCalledWhenServiceStarts(self):
+        """
+        The C{prepareConnection} callable is called after
+        L{ClientService.startService} once the connection is made.
+        """
+        prepares = [0]
+
+        def prepareConnection(_proto):
+            prepares[0] += 1
+
+        cq, service = self.makeReconnector(prepareConnection=prepareConnection,
+                                           startService=True)
+        self.assertEqual(1, prepares[0])
+
+
+    def test_prepareConnectionCalledWithProtocol(self):
+        """
+        The C{prepareConnection} callable is passed the connected protocol
+        instance.
+        """
+        newProtocols = []
+
+        def prepareConnection(proto):
+            newProtocols.append(proto)
+
+        cq, service = self.makeReconnector(
+            prepareConnection=prepareConnection,
+        )
+        self.assertIdentical(cq.constructedProtocols[0], newProtocols[0])
+
+
+    def test_prepareConnectionCalledAfterConnectionMade(self):
+        """
+        The C{prepareConnection} callback is invoked only once a connection is
+        made.
+        """
+        prepares = [0]
+
+        def prepareConnection(_proto):
+            prepares[0] += 1
+
+        clock = Clock()
+        cq, service = self.makeReconnector(prepareConnection=prepareConnection,
+                                           fireImmediately=False,
+                                           clock=clock)
+
+        cq.connectQueue[0].errback(Exception('connection attempt failed'))
+        self.assertEqual(0, prepares[0])  # Not called yet.
+
+        clock.advance(AT_LEAST_ONE_ATTEMPT)
+        cq.connectQueue[1].callback(None)
+
+        self.assertEqual(1, prepares[0])  # Was called.
+
+
+    def test_prepareConnectionCalledOnReconnect(self):
+        """
+        The C{prepareConnection} callback is invoked each time a connection is
+        made, including on reconnection.
+        """
+        prepares = [0]
+
+        def prepareConnection(_proto):
+            prepares[0] += 1
+
+        clock = Clock()
+        cq, service = self.makeReconnector(prepareConnection=prepareConnection,
+                                           clock=clock)
+
+        self.assertEqual(1, prepares[0])  # Called once.
+
+        # Protocol disconnects.
+        cq.constructedProtocols[0].connectionLost(Failure(IndentationError()))
+        clock.advance(AT_LEAST_ONE_ATTEMPT)
+
+        self.assertEqual(2, prepares[0])  # Called again.
+
+
+    def test_prepareConnectionReturnValueIgnored(self):
+        """
+        The C{prepareConnection} return value is ignored when it does not
+        indicate a failure. Even though the callback participates in the
+        internal new-connection L{Deferred} chain for error propagation
+        purposes, any successful result does not affect the ultimate return
+        value.
+        """
+        # Sentinel object returned by prepareConnection.
+        sentinel = object()
+
+        def prepareConnection(proto):
+            return sentinel
+
+        cq, service = self.makeReconnector(prepareConnection=prepareConnection)
+
+        result = self.successResultOf(service.whenConnected())
+        self.assertNotIdentical(sentinel, result)
+
+
+    def test_prepareConnectionReturningADeferred(self):
+        """
+        The C{prepareConnection} callable returns a deferred and calls to
+        L{ClientService.whenConnected} wait until it fires.
+        """
+        newProtocols = []
+        newProtocolDeferred = Deferred()
+
+        def prepareConnection(proto):
+            newProtocols.append(proto)
+            return newProtocolDeferred
+
+        cq, service = self.makeReconnector(prepareConnection=prepareConnection)
+
+        whenConnectedDeferred = service.whenConnected()
+        self.assertNoResult(whenConnectedDeferred)
+
+        newProtocolDeferred.callback(None)
+
+        self.assertIdentical(cq.applicationProtocols[0],
+                             self.successResultOf(whenConnectedDeferred))
+
+
+    def test_prepareConnectionThrows(self):
+        """
+        The connection attempt counts as a failure when the
+        C{prepareConnection} callable throws.
+        """
+        clock = Clock()
+
+        def prepareConnection(_proto):
+            raise IndentationError()
+
+        cq, service = self.makeReconnector(prepareConnection=prepareConnection,
+                                           clock=clock)
+
+        whenConnectedDeferred = service.whenConnected(failAfterFailures=2)
+        self.assertNoResult(whenConnectedDeferred)
+
+        clock.advance(AT_LEAST_ONE_ATTEMPT)
+        self.assertNoResult(whenConnectedDeferred)
+
+        clock.advance(AT_LEAST_ONE_ATTEMPT)
+        self.assertIdentical(IndentationError,
+                             self.failureResultOf(whenConnectedDeferred).type)
diff --git a/src/twisted/conch/scripts/ckeygen.py b/src/twisted/conch/scripts/ckeygen.py
@@ -189,7 +189,7 @@ def changePassPhrase(options):
             'Enter file in which the key is (%s): ' % filename)
     try:
         key = keys.Key.fromFile(options['filename'])
-    except keys.EncryptedKeyError as e:
+    except keys.EncryptedKeyError:
         # Raised if password not supplied for an encrypted key
         if not options.get('pass'):
             options['pass'] = getpass.getpass('Enter old passphrase: ')
diff --git a/src/twisted/internet/test/test_win32serialport.py b/src/twisted/internet/test/test_win32serialport.py
@@ -18,7 +18,7 @@
 try:
     from twisted.internet import serialport
     import serial
-except ImportError as e:
+except ImportError:
     serialport = None
     serial = None
 
diff --git a/src/twisted/newsfragments/8375.feature b/src/twisted/newsfragments/8375.feature
@@ -0,0 +1 @@
+twisted.application.internet.ClientService now accepts a function to initialize or validate a connection before it is returned by the whenConnected method as the prepareConnection argument.

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+twisted.application.internet.ClientService now accepts a function to initialize or validate a connection before it is returned by the whenConnected method as the prepareConnection argument.`