Updating system tests for Operation changes.

Also passing the current cluster when creating new operations.
dhermes · Dec 9, 2015 · 4ba9ae2 · 4ba9ae2
1 parent e6a49fc
commit 4ba9ae2
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 11 deletions.
diff --git a/gcloud_bigtable/cluster.py b/gcloud_bigtable/cluster.py
@@ -393,7 +393,7 @@ def create(self, timeout_seconds=None):
             request_pb, timeout_seconds)
 
         op_id, op_begin = _process_operation(cluster_pb.current_operation)
-        return Operation('create', op_id, op_begin)
+        return Operation('create', op_id, op_begin, cluster=self)
 
     def update(self, timeout_seconds=None):
         """Update this cluster.
@@ -430,7 +430,7 @@ def update(self, timeout_seconds=None):
             request_pb, timeout_seconds)
 
         op_id, op_begin = _process_operation(cluster_pb.current_operation)
-        return Operation('update', op_id, op_begin)
+        return Operation('update', op_id, op_begin, cluster=self)
 
     def delete(self, timeout_seconds=None):
         """Delete this cluster.
@@ -504,7 +504,7 @@ def undelete(self, timeout_seconds=None):
             request_pb, timeout_seconds)
 
         op_id, op_begin = _process_operation(operation_pb2)
-        return Operation('undelete', op_id, op_begin)
+        return Operation('undelete', op_id, op_begin, cluster=self)
 
     def list_tables(self, timeout_seconds=None):
         """List the tables in this cluster.

diff --git a/system_tests/run.py b/system_tests/run.py
@@ -120,13 +120,13 @@ def test_reload(self):
     def test_create_cluster(self):
         cluster_id = '%s-a' % (CLUSTER_ID,)
         cluster = CLIENT.cluster(CENTRAL_1C_ZONE, cluster_id)
-        cluster.create()
+        operation = cluster.create()
         # Make sure this cluster gets deleted after the test case.
         self.clusters_to_delete.append(cluster)
 
         # We want to make sure the operation completes.
         time.sleep(2)
-        self.assertTrue(cluster.operation_finished())
+        self.assertTrue(operation.finished())
 
         # Create a new cluster instance and make sure it is the same.
         cluster_alt = CLIENT.cluster(CENTRAL_1C_ZONE, cluster_id)
@@ -139,11 +139,11 @@ def test_create_cluster(self):
     def test_update(self):
         curr_display_name = CLUSTER.display_name
         CLUSTER.display_name = 'Foo Bar Baz'
-        CLUSTER.update()
+        operation = CLUSTER.update()
 
         # We want to make sure the operation completes.
         time.sleep(2)
-        self.assertTrue(CLUSTER.operation_finished())
+        self.assertTrue(operation.finished())
 
         # Create a new cluster instance and make sure it is the same.
         cluster_alt = CLIENT.cluster(CENTRAL_1C_ZONE, CLUSTER_ID)
@@ -156,11 +156,11 @@ def test_update(self):
         # Make sure to put the cluster back the way it was for the
         # other test cases.
         CLUSTER.display_name = curr_display_name
-        CLUSTER.update()
+        operation = CLUSTER.update()
 
         # We want to make sure the operation completes.
         time.sleep(2)
-        self.assertTrue(CLUSTER.operation_finished())
+        self.assertTrue(operation.finished())
 
 
 class TestTableAdminAPI(unittest2.TestCase):

diff --git a/system_tests/run_happybase.py b/system_tests/run_happybase.py
@@ -74,9 +74,9 @@ def set_cloud_bigtable_connection():
     try:
         cluster.reload()
     except NetworkError:
-        cluster.create()
+        operation = cluster.create()
         count = 0
-        while not cluster.operation_finished():
+        while not operation.finished():
             count += 1
             if count > 5:
                 raise NetworkError('Create cluster timed out')