Return an async_result from add_auth_async

Even though ZK closes the connection when auth fails (though the session is
still valid!) we should have add_auth_async be consistent with the rest
of the async calls by setting an exception on the AsyncResult object.

We still switch the state to KeeperState.AUTH_FAILED afterwards.

Added two tests (auth passes and auth fails raising AuthFailedError as
expected).

3 files changed, 39 insertions, 16 deletions

diff --git a/kazoo/client.py b/kazoo/client.py
index d506773..536864f 100644
--- a/kazoo/client.py
+++ b/kazoo/client.py
@@ -625,8 +625,16 @@ class KazooClient(object):
         :param scheme: authentication scheme (default supported:
                        "digest").
         :param credential: the credential -- value depends on scheme.
+
+        :returns: True if it was successful.
+        :rtype: bool
+
+        :raises:
+            :exc:`~kazoo.exceptions.AuthFailedError` if it failed though
+            the session state will be set to AUTH_FAILED as well.
+
         """
-        return self.add_auth_async(scheme, credential)
+        return self.add_auth_async(scheme, credential).get()
 
     def add_auth_async(self, scheme, credential):
         """Asynchronously send credentials to server. Takes the same
@@ -639,8 +647,10 @@ class KazooClient(object):
             raise TypeError("Invalid type for scheme")
         if not isinstance(credential, basestring):
             raise TypeError("Invalid type for credential")
-        self._call(Auth(0, scheme, credential), None)
-        return True
+
+        async_result = self.handler.async_result()
+        self._call(Auth(0, scheme, credential), async_result)
+        return async_result
 
     def unchroot(self, path):
         """Strip the chroot if applicable from the path."""
diff --git a/kazoo/protocol/connection.py b/kazoo/protocol/connection.py
index 38d9da3..ae11a67 100644
--- a/kazoo/protocol/connection.py
+++ b/kazoo/protocol/connection.py
@@ -387,13 +387,12 @@ class ConnectionHandler(object):
         elif header.xid == AUTH_XID:
             self.logger.debug('Received AUTH')
 
+            request, async_object, xid = client._pending.popleft()
             if header.err:
-                # We go ahead and fail out the connection, mainly because
-                # thats what Zookeeper client docs think is appropriate
-
-                # XXX TODO: Should we fail out? Or handle auth failure
-                # differently here since the session id is actually valid!
+                async_object.set_exception(AuthFailedError())
                 client._session_callback(KeeperState.AUTH_FAILED)
+            else:
+                async_object.set(True)
         elif header.xid == WATCH_XID:
             self._read_watch_event(buffer, offset)
         else:
@@ -427,17 +426,15 @@ class ConnectionHandler(object):
 
         # Special case for auth packets
         if request.type == Auth.type:
-            self._submit(request, connect_timeout, AUTH_XID)
-            client._queue.popleft()
-            os.read(self._read_pipe, 1)
-            return
-
-        self._xid += 1
+            xid = AUTH_XID
+        else:
+            self._xid += 1
+            xid = self._xid
 
-        self._submit(request, connect_timeout, self._xid)
+        self._submit(request, connect_timeout, xid)
         client._queue.popleft()
         os.read(self._read_pipe, 1)
-        client._pending.append((request, async_object, self._xid))
+        client._pending.append((request, async_object, xid))
 
     def _send_ping(self, connect_timeout):
         self.ping_outstanding.set()
diff --git a/kazoo/tests/test_client.py b/kazoo/tests/test_client.py
index 5b29ed9..47e7fce 100644
--- a/kazoo/tests/test_client.py
+++ b/kazoo/tests/test_client.py
@@ -193,6 +193,22 @@ class TestConnection(KazooTestCase):
         self.assertRaises(TypeError, self.client.add_auth,
             None, ('user', 'pass'))
 
+    def test_async_auth(self):
+        username = uuid.uuid4().hex
+        password = uuid.uuid4().hex
+        digest_auth = "%s:%s" % (username, password)
+        result = self.client.add_auth_async("digest", digest_auth)
+        self.assertTrue(result.get())
+
+    def test_async_auth_failure(self):
+        from kazoo.exceptions import AuthFailedError
+
+        username = uuid.uuid4().hex
+        password = uuid.uuid4().hex
+        digest_auth = "%s:%s" % (username, password)
+
+        self.assertRaises(AuthFailedError, self.client.add_auth, "unknown-scheme", digest_auth)
+
     def test_session_expire(self):
         from kazoo.protocol.states import KazooState