Add context manager for Bigtable client.

gcloud/bigtable/client.py

@@ -341,6 +341,10 @@ def start(self):
             self._operations_stub_internal.__enter__()
             self._table_stub_internal.__enter__()
 
+    def __enter__(self):
+        """Starts the client as a context manager."""
+        self.start()
+
     def stop(self):
         """Closes all the open gRPC clients."""
         if not self.is_started():
@@ -359,6 +363,10 @@ def stop(self):
         self._operations_stub_internal = None
         self._table_stub_internal = None
 
+    def __exit__(self, exc_type, exc_val, exc_t):
+        """Stops the client as a context manager."""
+        self.stop()
+
     def cluster(self, zone, cluster_id, display_name=None, serve_nodes=3):
         """Factory to create a cluster associated with this client.
 

gcloud/bigtable/test_client.py

@@ -110,6 +110,41 @@ def test_constructor_credentials_wo_create_scoped(self):
         expected_scopes = None
         self._constructor_test_helper(expected_scopes, creds)
 
+    def _context_manager_helper(self):
+        credentials = _Credentials()
+        project = 'PROJECT'
+        client = self._makeOne(project=project, credentials=credentials)
+
+        def mock_start():
+            client._data_stub_internal = object()
+        client.start = mock_start
+
+        def mock_stop():
+            client._data_stub_internal = None
+        client.stop = mock_stop
+        return client
+
+    def test_context_manager(self):
+        client = self._context_manager_helper()
+        self.assertFalse(client.is_started())
+        with client:
+            self.assertTrue(client.is_started())
+        self.assertFalse(client.is_started())
+
+    def test_context_manager_with_exception(self):
+        client = self._context_manager_helper()
+        self.assertFalse(client.is_started())
+
+        class DummyException(Exception):
+            pass
+        try:
+            with client:
+                self.assertTrue(client.is_started())
+                raise DummyException()
+        except DummyException:
+            pass
+        self.assertFalse(client.is_started())
+
     def _copy_test_helper(self, read_only=False, admin=False):
         credentials = _Credentials('value')
         project = 'PROJECT'