Adding _get_cluster helper to the HappyBase connection module.

This will replace client as the optional argument to
Connection since Cloud Bigtable accesses tables within
a cluster.

gcloud_bigtable/happybase/connection.py

@@ -32,6 +32,44 @@
 DEFAULT_PROTOCOL = None
 
 
+def _get_cluster(timeout=None):
+    """Gets cluster for the default project.
+
+    Creates a client with the inferred credentials and project ID from
+    the local environment. Then uses :meth:`.Client.list_clusters` to
+    get the unique cluster owned by the project.
+
+    If the request fails for any reason, or if there isn't exactly one cluster
+    owned by the project, then this function will fail.
+
+    :type timeout: int
+    :param timeout: (Optional) The socket timeout in milliseconds.
+
+    :rtype: :class:`gcloud_bigtable.cluster.Cluster`
+    :returns: The unique cluster owned by the project inferred from
+              the environment.
+    :raises: :class:`ValueError <exceptions.ValueError>` if any of the unused
+    """
+    client_kwargs = {'admin': True}
+    if timeout is not None:
+        client_kwargs['timeout_seconds'] = timeout / 1000.0
+    client = Client(**client_kwargs)
+    client.start()
+    clusters, failed_zones = client.list_clusters()
+    client.stop()
+
+    if len(failed_zones) != 0:
+        raise ValueError('Determining cluster via ListClusters encountered '
+                         'failed zones.')
+    if len(clusters) == 0:
+        raise ValueError('This client doesn\'t have access to any clusters.')
+    if len(clusters) > 1:
+        raise ValueError('This client has access to more than one cluster. '
+                         'Please directly pass the cluster you\'d '
+                         'like to use.')
+    return clusters[0]
+
+
 class Connection(object):
     """Connection to Cloud Bigtable backend.
 

gcloud_bigtable/happybase/test_connection.py

@@ -16,6 +16,60 @@
 import unittest2
 
 
+class Test__get_cluster(unittest2.TestCase):
+
+    def _callFUT(self, timeout=None):
+        from gcloud_bigtable.happybase.connection import _get_cluster
+        return _get_cluster(timeout=timeout)
+
+    def _helper(self, timeout=None, clusters=(), failed_zones=()):
+        from functools import partial
+        from gcloud_bigtable._testing import _Monkey
+        from gcloud_bigtable.happybase import connection as MUT
+
+        client_with_clusters = partial(_Client, clusters=clusters,
+                                       failed_zones=failed_zones)
+        with _Monkey(MUT, Client=client_with_clusters):
+            result = self._callFUT(timeout=timeout)
+
+        # If we've reached this point, then _callFUT didn't fail, so we know
+        # there is exactly one cluster.
+        cluster, = clusters
+        self.assertEqual(result, cluster)
+        client = cluster._client
+        self.assertEqual(client.args, ())
+        expected_kwargs = {'admin': True}
+        if timeout is not None:
+            expected_kwargs['timeout_seconds'] = timeout / 1000.0
+        self.assertEqual(client.kwargs, expected_kwargs)
+        self.assertEqual(client.start_calls, 1)
+        self.assertEqual(client.stop_calls, 1)
+
+    def test_default(self):
+        cluster = _Cluster()
+        self._helper(clusters=[cluster])
+
+    def test_with_timeout(self):
+        cluster = _Cluster()
+        self._helper(timeout=2103, clusters=[cluster])
+
+    def test_with_no_clusters(self):
+        with self.assertRaises(ValueError):
+            self._helper()
+
+    def test_with_too_many_clusters(self):
+        clusters = [_Cluster(), _Cluster()]
+        with self.assertRaises(ValueError):
+            self._helper(clusters=clusters)
+
+    def test_with_failed_zones(self):
+        cluster = _Cluster()
+        failed_zone = 'us-central1-c'
+        with self.assertRaises(ValueError):
+            self._helper(clusters=[cluster],
+                         failed_zones=[failed_zone])
+
+
 class TestConnection(unittest2.TestCase):
 
     def _getTargetClass(self):
@@ -40,19 +94,13 @@ def _constructor_missing_client_helper(self, timeout=None):
         from gcloud_bigtable._testing import _Monkey
         from gcloud_bigtable.happybase import connection as MUT
 
-        class FakeClient(object):
-
-            def __init__(self, *args, **kwargs):
-                self.args = args
-                self.kwargs = kwargs
-
-        with _Monkey(MUT, Client=FakeClient):
+        with _Monkey(MUT, Client=_Client):
             connection = self._makeOne(autoconnect=False, client=None,
                                        timeout=timeout)
             self.assertEqual(connection.table_prefix, None)
             self.assertEqual(connection.table_prefix_separator, '_')
             client = connection._client
-            self.assertTrue(isinstance(client, FakeClient))
+            self.assertTrue(isinstance(client, _Client))
             self.assertEqual(client.args, ())
             expected_kwargs = {'admin': True}
             if timeout is not None:
@@ -173,3 +221,30 @@ def test_compact_table(self):
         major = True
         with self.assertRaises(NotImplementedError):
             connection.compact_table(name, major=major)
+
+
+class _Client(object):
+
+    def __init__(self, *args, **kwargs):
+        self.clusters = kwargs.pop('clusters', [])
+        for cluster in self.clusters:
+            cluster._client = self
+        self.failed_zones = kwargs.pop('failed_zones', [])
+        self.args = args
+        self.kwargs = kwargs
+        self.start_calls = 0
+        self.stop_calls = 0
+
+    def start(self):
+        self.start_calls += 1
+
+    def stop(self):
+        self.stop_calls += 1
+
+    def list_clusters(self):
+        return self.clusters, self.failed_zones
+
+
+class _Cluster(object):
+
+    _client = None