Repository: smarter-travel-media/warthog

Issue:
Retry commands on requests ConnectionError and ConnectTimeout
These are pretty standard network-related nonsense that can cause issues during deploys and are safe to retry a few times. This issue came up when failing between members of a redundant LB pair.

Discussion:

Patch:
diff --git a/warthog/client.py b/warthog/client.py
index ff3d0fd..ec94955 100644
--- a/warthog/client.py
+++ b/warthog/client.py
@@ -120,7 +120,7 @@ class CommandFactory(object):
             self._transport_factory(), scheme_host, session_id, server)
 
 
-def _get_default_cmd_factory(verify, ssl_version):
+def _get_default_cmd_factory(verify, ssl_version, retries):
     """Get a :class:`CommandFactory` instance configured to use the provided TLS
     version and cert verification policy
 
@@ -129,39 +129,17 @@ def _get_default_cmd_factory(verify, ssl_version):
     :param int ssl_version: :mod:`ssl` module constant for specifying which SSL or
         TLS version to use for connecting to the load balancer over HTTPS, ``None``
         to use the default.
+    :param int retries: The maximum number of times to retry operations on transient
+        network errors.
     :return: Default command factory for building new commands to interact
         with the A10 load balancer.
     :rtype: WarthogCommandFactory
     """
     return CommandFactory(warthog.transport.get_transport_factory(
-        verify=verify, ssl_version=ssl_version
+        verify=verify, ssl_version=ssl_version, retries=retries
     ))
 
 
-@contextlib.contextmanager
-def session_context(scheme_host, username, password, commands):
-    """Context manager that makes a request to start an authenticated session, yields the
-    session ID, and then closes the session afterwards.
-
-    :param basestring scheme_host: Scheme, host, and port combination of the load balancer.
-    :param basestring username: Name of the user to authenticate with.
-    :param basestring password: Password for the user to authenticate with.
-    :param CommandFactory commands: Factory instance for creating new commands
-        for starting and ending sessions with the load balancer.
-    :return: The session ID of the newly established session.
-    """
-    session = None
-    try:
-        start_cmd = commands.get_session_start(scheme_host, username, password)
-        session = start_cmd.send()
-
-        yield session
-    finally:
-        if session is not None:
-            end_cmd = commands.get_session_end(scheme_host, session)
-            end_cmd.send()
-
-
 class WarthogClient(object):
     """Client for interacting with an A10 load balancer to get the status
     of nodes managed by it, enable them, and disable them.
@@ -172,27 +150,25 @@ class WarthogClient(object):
         Removed .disabled_context() method.
     """
     _logger = warthog.core.get_log()
-    _default_wait_interval = 2.0
 
     # pylint: disable=too-many-arguments
     def __init__(self, scheme_host, username, password,
                  verify=None,
                  ssl_version=None,
-                 wait_interval=_default_wait_interval,
+                 network_retries=None,
                  commands=None):
         """Set the load balancer scheme/host/port combination, username and password
         to use for connecting and authenticating with the load balancer.
 
-        Optionally, whether or not to verify certificates when using HTTPS may be
-        toggled. This can enable you to use a self signed certificate for the load
-        balancer while still using HTTPS.
+        Whether or not to verify certificates when using HTTPS may be toggled via the
+        ``verify`` parameter. This can enable you to use a self signed certificate for
+        the load balancer while still using HTTPS.
 
-        Optionally, the version of SSL or TLS to use may be specified as a :mod:`ssl`
-        module protocol constant.
+        The version of SSL or TLS to use may be specified as a :mod:`ssl` module protocol
+        constant via the ``ssl_version`` parameter.
 
-        Optionally, the amount of time to wait between retries of various operations
-        and the factory used for creating commands may be set. If the interval between
-        retries is not supplied, the default is two seconds.
+        The maximum number of times to retry network operations on transient errors
+        can be specified via the ``network_retries`` parameter.
 
         If the command factory is not supplied, a default instance will be used. The
         command factory is responsible for creating new :class:`requests.Session` instances
@@ -207,6 +183,16 @@ class WarthogClient(object):
             Added the optional ``ssl_version`` parameter to make use of alternate SSL
             or TLS versions easier.
 
+        .. versionchanged:: 2.0.0
+            Added the optional ``network_retries`` parameter to make use of retry logic on
+            transient network errors. A non-zero number of retries is used by default if
+            this is not specified. Previously, no retries were attempted on transient network
+            errors.
+
+        .. versionchanged:: 2.0.0
+            Removed the optional ``wait_interval`` parameter. This is now passed directly
+            as an argument to :meth:`enable_node` or :meth:`disable_node` methods.
+
         :param basestring scheme_host: Scheme, host, and port combination of the load balancer.
         :param basestring username: Name of the user to authenticate with.
         :param basestring password: Password for the user to authenticate with.
@@ -216,23 +202,36 @@ class WarthogClient(object):
         :param int|None ssl_version: :mod:`ssl` module constant for specifying which version of
             SSL or TLS to use when connecting to the load balancer over HTTPS, ``None`` to use
             the library default. The default is to use TLSv1.2.
-        :param float wait_interval: How long (in seconds) to wait between each retry of
-            various operations (waiting for nodes to transition, waiting for connections
-            to close, etc.).
+        :param int|None network_retries: Maximum number of times to retry network operations on
+            transient network errors. Default is to retry network operations a non-zero number
+            of times.
         :param CommandFactory commands: Factory instance for creating new commands for
             starting and ending sessions with the load balancer.
         """
         self._scheme_host = scheme_host
         self._username = username
         self._password = password
-        self._interval = wait_interval
         self._commands = commands if commands is not None else \
-            _get_default_cmd_factory(verify, ssl_version)
+            _get_default_cmd_factory(verify, ssl_version, network_retries)
 
+    @contextlib.contextmanager
     def _session_context(self):
-        """Get a new context manager that starts and ends a session with the load balancer."""
+        """Context manager that makes a request to start an authenticated session, yields the
+        session ID, and then closes the session afterwards.
+
+        :return: The session ID of the newly established session.
+        """
         self._logger.debug('Creating new session context for %s', self._scheme_host)
-        return session_context(self._scheme_host, self._username, self._password, self._commands)
+        session = None
+        try:
+            start_cmd = self._commands.get_session_start(self._scheme_host, self._username, self._password)
+            session = start_cmd.send()
+
+            yield session
+        finally:
+            if session is not None:
+                end_cmd = self._commands.get_session_end(self._scheme_host, session)
+                end_cmd.send()
 
     def get_status(self, server):
         """Get the current status of the given server, at the node level.
@@ -248,12 +247,11 @@ class WarthogClient(object):
             operation.
         :raises warthog.exceptions.WarthogNoSuchNodeError: If the load balancer does
             not recognize the given hostname.
-        :raises warthog.exceptions.WarthogNodeStatusError: If there are any other
+        :raises warthog.exceptions.WarthogApiError: If there are any other
             problems getting the status of the given server.
         """
         with self._session_context() as session:
-            cmd = self._commands.get_server_status(
-                self._scheme_host, session, server)
+            cmd = self._commands.get_server_status(self._scheme_host, session, server)
             return cmd.send()
 
     def get_connections(self, server):
@@ -271,28 +269,31 @@ class WarthogClient(object):
             operation.
         :raises warthog.exceptions.WarthogNoSuchNodeError: If the load balancer does
             not recognize the given hostname.
-        :raises warthog.exceptions.WarthogNodeStatusError: If there are any other
+        :raises warthog.exceptions.WarthogApiError: If there are any other
             problems getting the active connections for the given server.
 
         .. versionadded:: 0.4.0
         """
         with self._session_context() as session:
-            cmd = self._commands.get_active_connections(
-                self._scheme_host, session, server)
+            cmd = self._commands.get_active_connections(self._scheme_host, session, server)
             return cmd.send()
 
-    def disable_server(self, server, max_retries=5):
+    def disable_server(self, server, max_retries=5, wait_interval=2.0):
         """Disable a server at the node level, optionally retrying when there are transient
         errors and waiting for the number of active connections to the server to reach zero.
 
-        If ``max_retries`` is zero, no attempt will be made to retry on transient errors
-        or to wait until there are no active connections to the server, the method will
-        try a single time to disable the server and then return immediately.
+        If ``max_retries`` is zero, no attempt will be made to wait until there are no active
+        connections to the server, the method will try a single time to disable the server and
+        then return immediately.
+
+        .. versionchanged:: 2.0.0
+            Added the optional ``wait_interval`` parameter.
 
         :param basestring server: Hostname of the server to disable
-        :param int max_retries: Max number of times to sleep and retry when encountering
-            some sort of transient error when disabling the server and while waiting for
+        :param int max_retries: Max number of times to sleep and retry while waiting for
             the number of active connections to a server to reach zero.
+        :param float wait_interval: How long (in seconds) to wait between each check to
+            see if the number of active connections to a server has reached zero.
         :return: True if the server was disabled, false otherwise.
         :rtype: bool
         :raises warthog.exceptions.WarthogAuthFailureError: If authentication with
@@ -300,26 +301,23 @@ class WarthogClient(object):
             operation.
         :raises warthog.exceptions.WarthogNoSuchNodeError: If the load balancer does
             not recognize the given hostname.
-        :raises warthog.exceptions.WarthogNodeDisableError: If there are any other
+        :raises warthog.exceptions.WarthogApiError: If there are any other
             problems disabling the given server.
         """
         with self._session_context() as session:
             disable = self._commands.get_disable_server(self._scheme_host, session, server)
-            self._try_repeatedly(disable.send, max_retries)
+            disable.send()
 
             active = self._commands.get_active_connections(self._scheme_host, session, server)
-            self._wait_for_connections(active.send, max_retries)
+            self._wait_for_connections(active.send, max_retries, wait_interval)
 
-            status = self._commands.get_server_status(
-                self._scheme_host, session, server)
+            status = self._commands.get_server_status(self._scheme_host, session, server)
             return warthog.core.STATUS_DISABLED == status.send()
 
-    # NOTE: there's a fair amount of duplicate code between this method and _wait_for_status
-    # and we could consolidate them to one method that just accepts a function and waits for
-    # it to return true and then break. But, this way we have more useful debug information
-    # logged at the expense of duplicate code.
-    # pylint: disable=missing-docstring
-    def _wait_for_connections(self, conn_method, max_retries):
+    def _wait_for_connections(self, conn_method, max_retries, interval):
+        """Repeatedly execute a command to get the number of active connections until
+        the number of active connections drops to zero or we run out of retries.
+        """
         retries = 0
 
         while retries < max_retries:
@@ -328,21 +326,26 @@ class WarthogClient(object):
                 break
 
             self._logger.debug(
-                "Connections still active: %s, sleeping for %s seconds...", conns, self._interval)
-            time.sleep(self._interval)
+                "Connections still active: %s, sleeping for %s seconds...", conns, interval)
+            time.sleep(interval)
             retries += 1
 
-    def enable_server(self, server, max_retries=5):
+    def enable_server(self, server, max_retries=5, wait_interval=2.0):
         """Enable a server at the node level, optionally retrying when there are transient
         errors and waiting for the server to enter the expected, enabled state.
 
-        If ``max_retries`` is zero, no attempt will be made to retry on transient errors
-        or to wait until the server enters the expected, enabled state, the method will
-        try a single time to enable the server then return immediately.
+        If ``max_retries`` is zero, no attempt will be made to wait until the server enters
+        the expected, enabled state, the method will try a single time to enable the server
+        then return immediately.
+
+        .. versionchanged:: 2.0.0
+            Added the optional ``wait_interval`` parameter.
 
         :param basestring server: Hostname of the server to enable
-        :param int max_retries: Max number of times to sleep and retry when encountering
-            some transient error while trying to enable the server
+        :param int max_retries: Max number of times to sleep and retry while waiting for
+            the server to enter the "enabled" state.
+        :param float wait_interval: How long (in seconds) to wait between each check to
+            see if the server has entered the "enabled" state.
         :return: True if the server was enabled, false otherwise
         :rtype: bool
         :raises warthog.exceptions.WarthogAuthFailureError: If authentication with
@@ -350,20 +353,22 @@ class WarthogClient(object):
             operation.
         :raises warthog.exceptions.WarthogNoSuchNodeError: If the load balancer does
             not recognize the given hostname.
-        :raises warthog.exceptions.WarthogNodeEnableError: If there are any other
+        :raises warthog.exceptions.WarthogApiError: If there are any other
             problems enabling the given server.
         """
         with self._session_context() as session:
             enable = self._commands.get_enable_server(self._scheme_host, session, server)
-            self._try_repeatedly(enable.send, max_retries)
+            enable.send()
 
             status = self._commands.get_server_status(self._scheme_host, session, server)
-            self._wait_for_enable(status.send, max_retries)
+            self._wait_for_enable(status.send, max_retries, wait_interval)
 
             return warthog.core.STATUS_ENABLED == status.send()
 
-    # pylint: disable=missing-docstring
-    def _wait_for_enable(self, status_method, max_retries):
+    def _wait_for_enable(self, status_method, max_retries, interval):
+        """Repeatedly execute a command to get the status of a node until the node
+        becomes enabled or we run out of retries.
+        """
         retries = 0
 
         while retries < max_retries:
@@ -373,28 +378,6 @@ class WarthogClient(object):
 
             self._logger.debug(
                 "Server is not yet enabled (%s), sleeping for %s seconds...",
-                status, self._interval)
-            time.sleep(self._interval)
+                status, interval)
+            time.sleep(interval)
             retries += 1
-
-    def _try_repeatedly(self, method, max_retries):
-        """Execute a method, retrying if it fails due to a transient error
-        up to a given number of times, with the instance-wide interval in
-        between each try.
-        """
-        retries = 0
-
-        while True:
-            try:
-                return method()
-            except warthog.exceptions.WarthogApiError as e:
-                if e.api_code not in warthog.core.TRANSIENT_ERRORS or retries >= max_retries:
-                    raise
-                self._logger.debug(
-                    "Encountered transient error %s - %s, retrying... ", e.api_code, e.api_msg)
-                time.sleep(self._interval)
-                retries += 1
-
-
-# NOTE: This alias is only for transitioning to the v3 API
-CommandFactory3 = CommandFactory
diff --git a/warthog/core.py b/warthog/core.py
index 147d1f1..912f48b 100644
--- a/warthog/core.py
+++ b/warthog/core.py
@@ -32,8 +32,6 @@ ERROR_CODE_NO_SUCH_SERVER = 1023460352
 
 ERROR_CODE_BAD_PERMISSION = 419545856
 
-TRANSIENT_ERRORS = frozenset([])
-
 _PATH_AUTH = '/axapi/v3/auth'
 
 _PATH_LOGOFF = '/axapi/v3/logoff'
diff --git a/warthog/transport.py b/warthog/transport.py
index 776d2c8..84134fb 100644
--- a/warthog/transport.py
+++ b/warthog/transport.py
@@ -21,8 +21,7 @@ import requests
 from requests.adapters import (
     HTTPAdapter,
     DEFAULT_POOLBLOCK,
-    DEFAULT_POOLSIZE,
-    DEFAULT_RETRIES)
+    DEFAULT_POOLSIZE)
 from requests.packages.urllib3.exceptions import InsecureRequestWarning
 from requests.packages.urllib3.poolmanager import PoolManager
 
@@ -36,28 +35,34 @@ DEFAULT_SSL_VERSION = warthog.ssl.PROTOCOL_TLSv1_2
 # Default to verifying SSL/TLS certs because "safe by default" is a good idea.
 DEFAULT_CERT_VERIFY = True
 
+# Default number of times to retry on transient network errors like connection
+# timeouts or DNS timeouts.
+DEFAULT_RETRIES = 5
 
-def get_transport_factory(verify=None, ssl_version=None):
+
+def get_transport_factory(verify=None, ssl_version=None, retries=None):
     """Get a new callable that returns :class:`requests.Session` instances that
     have been configured according to the given parameters.
 
     :class:`requests.Session` instances are then used for interacting with the API
     of the load balancer over HTTP or HTTPS.
 
-    It is typically not required for user code to call this function directly unless
-    you have special requirements such as needing to bypass HTTPS certificate validation
-    because you use a self signed certificate.
-
     .. versionchanged:: 0.10.0
         Using the requests/urllib3 default is no longer an option. Passing a ``None`` value
         for ``ssl_version`` will result in using the Warthog default (TLS v1).
 
+    .. versionchanged:: 2.0.0
+        Added the ``retries`` parameter and default it to a number greater than zero.
+
     :param bool|None verify: Should SSL certificates by verified when connecting
         over HTTPS? Default is ``True``. If you have chosen not to verify certificates
         warnings about this emitted by the requests library will be suppressed.
     :param int|None ssl_version: Explicit version of SSL to use for HTTPS connections
         to an A10 load balancer. The version is a constant as specified by the
         :mod:`ssl` module. The default is TLSv1.
+    :param int|None retries: The maximum number of times to retry operations on transient
+        network errors. Note this only applies to cases where we haven't yet sent any
+        data to the server (e.g. connection errors, DNS errors, etc.)
     :return: A callable to return new configured session instances for making HTTP(S)
         requests
     :rtype: callable
@@ -67,15 +72,23 @@ def get_transport_factory(verify=None, ssl_version=None):
     # just pass `None` and we'll pick the default here.
     verify = verify if verify is not None else DEFAULT_CERT_VERIFY
     ssl_version = ssl_version if ssl_version is not None else DEFAULT_SSL_VERSION
+    retries = retries if retries is not None else DEFAULT_RETRIES
 
     # pylint: disable=missing-docstring
     def factory():
         transport = requests.Session()
-        transport.mount('https://', VersionedSSLAdapter(ssl_version))
+        transport.mount('https://', VersionedSSLAdapter(ssl_version, max_retries=retries))
 
         if not verify:
             transport.verify = False
 
+        transport.mount('http://', HTTPAdapter(
+            max_retries=retries,
+            pool_connections=DEFAULT_POOLSIZE,
+            pool_maxsize=DEFAULT_POOLSIZE,
+            pool_block=DEFAULT_POOLBLOCK
+        ))
+
         return transport
 
     # Make sure that we suppress warnings about invalid certs since the user
@@ -100,7 +113,8 @@ class VersionedSSLAdapter(HTTPAdapter):
             pool_connections=pool_connections,
             pool_maxsize=pool_maxsize,
             max_retries=max_retries,
-            pool_block=pool_block)
+            pool_block=pool_block
+        )
 
     def init_poolmanager(self, connections, maxsize, block=False, **pool_kwargs):
         # pylint: disable=attribute-defined-outside-init


Test patch:
diff --git a/test/test_client.py b/test/test_client.py
index 342f9a5..2c80aee 100644
--- a/test/test_client.py
+++ b/test/test_client.py
@@ -1,9 +1,8 @@
 # -*- coding: utf-8 -*-
 
-import logging
-
-import pytest
 import mock
+import pytest
+
 import warthog.client
 import warthog.core
 import warthog.exceptions
@@ -59,15 +58,13 @@ def disable_cmd():
     return mock.Mock(spec=warthog.core.NodeDisableCommand)
 
 
-@pytest.fixture
-def client():
-    return mock.Mock(spec=warthog.client.WarthogClient)
-
-
 def test_session_context_enter_yields_session(commands, start_cmd):
     start_cmd.send.return_value = '1234'
 
-    context = warthog.client.session_context(SCHEME_HOST, 'user', 'password', commands)
+    client = warthog.client.WarthogClient(
+        SCHEME_HOST, 'user', 'password', commands=commands)
+
+    context = client._session_context()
     session = context.__enter__()
 
     assert '1234' == session, 'Did not get expected session ID'
@@ -78,7 +75,10 @@ def test_session_context_enter_yields_session(commands, start_cmd):
 def test_session_context_exit_closes_previous_session(commands, start_cmd, end_cmd):
     start_cmd.send.return_value = '1234'
 
-    context = warthog.client.session_context(SCHEME_HOST, 'user', 'password', commands)
+    client = warthog.client.WarthogClient(
+        SCHEME_HOST, 'user', 'password', commands=commands)
+
+    context = client._session_context()
     context.__enter__()
     context.__exit__(None, None, None)
 
@@ -88,92 +88,100 @@ def test_session_context_exit_closes_previous_session(commands, start_cmd, end_c
 
 def test_session_context_exception_in_context_propagated(commands, start_cmd):
     start_cmd.send.return_value = '1234'
-    context = warthog.client.session_context(SCHEME_HOST, 'user', 'password', commands)
+
+    client = warthog.client.WarthogClient(
+        SCHEME_HOST, 'user', 'password', commands=commands)
+
+    context = client._session_context()
 
     with pytest.raises(RuntimeError):
         with context:
             raise RuntimeError("AHH!")
 
 
-class TestWarthogClient(object):
-    def test_get_status(self, commands, start_cmd, end_cmd, status_cmd):
-        start_cmd.send.return_value = '1234'
-        status_cmd.send.return_value = 'down'
-
-        client = warthog.client.WarthogClient(
-            SCHEME_HOST, 'user', 'password', wait_interval=0.1, commands=commands)
+def test_get_status(commands, start_cmd, end_cmd, status_cmd):
+    start_cmd.send.return_value = '1234'
+    status_cmd.send.return_value = 'down'
 
-        status = client.get_status('app1.example.com')
+    client = warthog.client.WarthogClient(
+        SCHEME_HOST, 'user', 'password', commands=commands)
 
-        assert 'down' == status, 'Did not get expected status'
-        assert end_cmd.send.called, 'Session end .send() did not get called'
+    status = client.get_status('app1.example.com')
 
-    def test_get_connections(self, commands, start_cmd, end_cmd, conn_cmd):
-        start_cmd.send.return_value = '1234'
-        conn_cmd.send.return_value = 42
+    assert 'down' == status, 'Did not get expected status'
+    assert end_cmd.send.called, 'Session end .send() did not get called'
 
-        client = warthog.client.WarthogClient(
-            SCHEME_HOST, 'user', 'password', wait_interval=0.1, commands=commands)
 
-        connections = client.get_connections('app1.example.com')
+def test_get_connections(commands, start_cmd, end_cmd, conn_cmd):
+    start_cmd.send.return_value = '1234'
+    conn_cmd.send.return_value = 42
 
-        assert 42 == connections, 'Did not get expected active connections'
-        assert end_cmd.send.called, 'Session end .send() did not get called'
+    client = warthog.client.WarthogClient(
+        SCHEME_HOST, 'user', 'password', commands=commands)
 
-    def test_disable_server_no_active_connections(self, commands, start_cmd, end_cmd,
-                                                  status_cmd, conn_cmd, disable_cmd):
-        start_cmd.send.return_value = '1234'
-        disable_cmd.send.return_value = True
-        conn_cmd.send.return_value = 0
-        status_cmd.send.return_value = 'disabled'
+    connections = client.get_connections('app1.example.com')
 
-        client = warthog.client.WarthogClient(
-            SCHEME_HOST, 'user', 'password', wait_interval=0.1, commands=commands)
+    assert 42 == connections, 'Did not get expected active connections'
+    assert end_cmd.send.called, 'Session end .send() did not get called'
 
-        disabled = client.disable_server('app1.example.com')
 
-        assert disabled, 'Server did not end up disabled'
-        assert end_cmd.send.called, 'Session end .send() did not get called'
+def test_disable_server_no_active_connections(commands, start_cmd, end_cmd,
+                                              status_cmd, conn_cmd, disable_cmd):
+    start_cmd.send.return_value = '1234'
+    disable_cmd.send.return_value = True
+    conn_cmd.send.return_value = 0
+    status_cmd.send.return_value = 'disabled'
 
-    def test_disable_server_with_active_connections(self, commands, start_cmd, end_cmd,
-                                                    status_cmd, conn_cmd, disable_cmd):
-        start_cmd.send.return_value = '1234'
-        disable_cmd.send.return_value = True
-        conn_cmd.send.return_value = [42, 3, 0]
-        status_cmd.send.return_value = 'disabled'
+    client = warthog.client.WarthogClient(
+        SCHEME_HOST, 'user', 'password', commands=commands)
 
-        client = warthog.client.WarthogClient(
-            SCHEME_HOST, 'user', 'password', wait_interval=0.1, commands=commands)
+    disabled = client.disable_server('app1.example.com')
 
-        disabled = client.disable_server('app1.example.com')
+    assert disabled, 'Server did not end up disabled'
+    assert end_cmd.send.called, 'Session end .send() did not get called'
 
-        assert disabled, 'Server did not end up disabled'
-        assert end_cmd.send.called, 'Session end .send() did not get called'
 
-    def test_disable_server_never_gets_disabled(self, commands, start_cmd, end_cmd,
+def test_disable_server_with_active_connections(commands, start_cmd, end_cmd,
                                                 status_cmd, conn_cmd, disable_cmd):
-        start_cmd.send.return_value = '1234'
-        disable_cmd.send.return_value = True
-        conn_cmd.send.return_value = 42
-        status_cmd.send.return_value = 'enabled'
+    start_cmd.send.return_value = '1234'
+    disable_cmd.send.return_value = True
+    conn_cmd.send.return_value = [42, 3, 0]
+    status_cmd.send.return_value = 'disabled'
+
+    client = warthog.client.WarthogClient(
+        SCHEME_HOST, 'user', 'password', commands=commands)
+
+    disabled = client.disable_server('app1.example.com', wait_interval=0.1)
 
-        client = warthog.client.WarthogClient(
-            SCHEME_HOST, 'user', 'password', wait_interval=0.1, commands=commands)
+    assert disabled, 'Server did not end up disabled'
+    assert end_cmd.send.called, 'Session end .send() did not get called'
 
-        disabled = client.disable_server('app1.example.com')
 
-        assert not disabled, 'Server ended up disabled'
-        assert end_cmd.send.called, 'Session end .send() did not get called'
+def test_disable_server_never_gets_disabled(commands, start_cmd, end_cmd,
+                                            status_cmd, conn_cmd, disable_cmd):
+    start_cmd.send.return_value = '1234'
+    disable_cmd.send.return_value = True
+    conn_cmd.send.return_value = 42
+    status_cmd.send.return_value = 'enabled'
+
+    client = warthog.client.WarthogClient(
+        SCHEME_HOST, 'user', 'password', commands=commands)
+
+    disabled = client.disable_server('app1.example.com', wait_interval=0.1)
 
-    def test_enable_server(self, commands, start_cmd, end_cmd, status_cmd, enable_cmd):
-        start_cmd.send.return_value = '1234'
-        enable_cmd.send.return_value = True
-        status_cmd.send.side_effect = ['down', 'down', 'enabled', 'enabled']
+    assert not disabled, 'Server ended up disabled'
+    assert end_cmd.send.called, 'Session end .send() did not get called'
+
+
+def test_enable_server(commands, start_cmd, end_cmd, status_cmd, enable_cmd):
+    start_cmd.send.return_value = '1234'
+    enable_cmd.send.return_value = True
+    status_cmd.send.side_effect = ['down', 'down', 'enabled', 'enabled']
 
-        client = warthog.client.WarthogClient(
-            SCHEME_HOST, 'user', 'password', wait_interval=0.1, commands=commands)
+    client = warthog.client.WarthogClient(
+        SCHEME_HOST, 'user', 'password', commands=commands)
 
-        enabled = client.enable_server('app1.example.com')
+    enabled = client.enable_server('app1.example.com', wait_interval=0.1)
 
-        assert enabled, 'Server did not end up enabled'
-        assert end_cmd.send.called, 'Session end .send() did not get called'
+    assert enabled, 'Server did not end up enabled'
+    assert end_cmd.send.called, 'Session end .send() did not get called'