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

@@ -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

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])

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

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()