From 26807891bfbf29cee3eddb661f8233ea42492107 Mon Sep 17 00:00:00 2001
From: cyli <cyli@twistedmatrix.com>
Date: Mon, 15 Jul 2013 10:24:02 -0700
Subject: [PATCH] Apply @bmuller's changes from pull request #18 - move the
 disconnect method from TestCQLClient to CQLClient and add a disconnect method
 and tests to cluster.

---
 silverberg/client.py            | 17 ++++++++---------
 silverberg/cluster.py           | 12 ++++++++++++
 silverberg/test/test_client.py  | 18 +++++++++---------
 silverberg/test/test_cluster.py |  7 +++++++
 4 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/silverberg/client.py b/silverberg/client.py
index 02532e6..1b1e902 100644
--- a/silverberg/client.py
+++ b/silverberg/client.py
@@ -86,6 +86,14 @@ def _handshake(client):
         ds = self._client.connection(_handshake)
         return ds
 
+    def disconnect(self):
+        """
+        Disconnect from the cassandra cluster.
+
+        :return: a :class:`Deferred` that fires with None when disconnected.
+        """
+        return self._client.disconnect()
+
     def describe_version(self):
         """
         Query the Cassandra server for the version.
@@ -210,15 +218,6 @@ def transport(self):
         """
         return self._client._transport
 
-    def disconnect(self):
-        """
-        Disconnect from the cassandra cluster.  Likely to be used for testing
-        purposes only.
-
-        :return: a :class:`Deferred` that fires with None when disconnected.
-        """
-        return self._client.disconnect()
-
     def pause(self):
         """
         Pause the client by removing the connection from the reactor.  This is
diff --git a/silverberg/cluster.py b/silverberg/cluster.py
index 92f81e0..b9a71f8 100644
--- a/silverberg/cluster.py
+++ b/silverberg/cluster.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from twisted.internet.defer import DeferredList
+
 from silverberg.client import CQLClient
 
 
@@ -44,3 +46,13 @@ def execute(self, *args, **kwargs):
         See :py:func:`silverberg.client.CQLClient.execute`
         """
         return self._client().execute(*args, **kwargs)
+
+    def disconnect(self):
+        """
+        Disconnect every client from the cassandra cluster.  Likely to be used for testing
+        purposes only.
+
+        :return: a :class:`DeferredList` that fires with a list of None's when every client
+        has disconnected.
+        """
+        return DeferredList([client.disconnect() for client in self._seed_clients])
diff --git a/silverberg/test/test_client.py b/silverberg/test/test_client.py
index ae080d1..2420748 100644
--- a/silverberg/test/test_client.py
+++ b/silverberg/test/test_client.py
@@ -54,6 +54,15 @@ def _connect(factory):
 
         self.endpoint.connect.side_effect = _connect
 
+    def test_disconnect(self):
+        """
+        When disconnect is called, the on demand thrift client is disconnected
+        """
+        client = TestingCQLClient(self.endpoint, 'blah')
+        self.assertFired(client.describe_version())
+        client.disconnect()
+        self.twisted_transport.loseConnection.assert_called_once_with()
+
     def test_login(self):
         """Test that login works as expected."""
         client = CQLClient(self.endpoint, 'blah', 'groucho', 'swordfish')
@@ -259,15 +268,6 @@ def test_transport_exposed(self):
         self.assertFired(client.describe_version())
         self.assertIs(client.transport, self.twisted_transport)
 
-    def test_disconnect(self):
-        """
-        When disconnect is called, the on demand thrift client is disconnected
-        """
-        client = TestingCQLClient(self.endpoint, 'blah')
-        self.assertFired(client.describe_version())
-        client.disconnect()
-        self.twisted_transport.loseConnection.assert_called_once_with()
-
     def test_pause(self):
         """
         When pausing, stop reading and stop writing on the transport are called
diff --git a/silverberg/test/test_cluster.py b/silverberg/test/test_cluster.py
index 9947c45..af310a1 100644
--- a/silverberg/test/test_cluster.py
+++ b/silverberg/test/test_cluster.py
@@ -43,3 +43,10 @@ def test_round_robin_execute(self):
             result = cluster.execute(arg)
             self.clients[client].execute.assert_called_with(arg)
             self.assertEqual(self.clients[client].execute.return_value, result)
+
+    def test_disconnect(self):
+        cluster = RoundRobinCassandraCluster(['one', 'two', 'three'], 'keyspace')
+        cluster.disconnect()
+
+        for client in self.clients:
+            client.disconnect.assert_called_with()