Making sure zones are in OK state in list_zones.

dhermes · Jul 28, 2015 · 1638be0 · 1638be0
1 parent 6c0ae72
commit 1638be0
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 4 deletions.
diff --git a/gcloud_bigtable/client.py b/gcloud_bigtable/client.py
@@ -36,6 +36,7 @@
 except ImportError:
     app_identity = None
 
+from gcloud_bigtable._generated import bigtable_cluster_data_pb2 as data_pb2
 from gcloud_bigtable._generated import (
     bigtable_cluster_service_messages_pb2 as messages_pb2)
 from gcloud_bigtable.cluster import Cluster
@@ -343,6 +344,8 @@ def list_zones(self, timeout_seconds=TIMEOUT_SECONDS):
 
         :rtype: list of strings
         :returns: The names of the zones
+        :raises: :class:`ValueError` if one of the zones is not in
+                 ``OK`` state.
         """
         request_pb = messages_pb2.ListZonesRequest(name=self.project_name)
         stub = make_stub(self._credentials, CLUSTER_STUB_FACTORY,
@@ -352,7 +355,13 @@ def list_zones(self, timeout_seconds=TIMEOUT_SECONDS):
             # We expect a `messages_pb2.ListZonesResponse`
             list_zones_response = response.result()
 
-        return [zone.display_name for zone in list_zones_response.zones]
+        result = []
+        for zone in list_zones_response.zones:
+            if zone.status != data_pb2.Zone.OK:
+                raise ValueError('Zone %s not in OK state' % (
+                    zone.display_name,))
+            result.append(zone.display_name)
+        return result
 
     def list_clusters(self, timeout_seconds=TIMEOUT_SECONDS):
         """Lists clusters owned by the project.

diff --git a/gcloud_bigtable/test_client.py b/gcloud_bigtable/test_client.py
@@ -421,7 +421,7 @@ def test_cluster_factory(self):
         self.assertEqual(cluster.zone, zone)
         self.assertEqual(cluster.cluster_id, cluster_id)
 
-    def test_list_zones(self):
+    def _list_zones_helper(self, zone_status):
         from gcloud_bigtable._generated import (
             bigtable_cluster_data_pb2 as data_pb2)
         from gcloud_bigtable._generated import (
@@ -434,8 +434,8 @@ def test_list_zones(self):
         zone2 = 'foo'
         response_pb = messages_pb2.ListZonesResponse(
             zones=[
-                data_pb2.Zone(display_name=zone1),
-                data_pb2.Zone(display_name=zone2),
+                data_pb2.Zone(display_name=zone1, status=zone_status),
+                data_pb2.Zone(display_name=zone2, status=zone_status),
             ],
         )
         expected_result = [zone1, zone2]
@@ -446,6 +446,17 @@ def result_method(client):
         self._grpc_client_test_helper('ListZones', result_method, request_pb,
                                       response_pb, expected_result, PROJECT_ID)
 
+    def test_list_zones(self):
+        from gcloud_bigtable._generated import (
+            bigtable_cluster_data_pb2 as data_pb2)
+        self._list_zones_helper(data_pb2.Zone.OK)
+
+    def test_list_zones_failure(self):
+        from gcloud_bigtable._generated import (
+            bigtable_cluster_data_pb2 as data_pb2)
+        with self.assertRaises(ValueError):
+            self._list_zones_helper(data_pb2.Zone.EMERGENCY_MAINENANCE)
+
     def test_list_clusters(self):
         from gcloud_bigtable._generated import (
             bigtable_cluster_data_pb2 as data_pb2)