From 87aed7f6d8efac5d7a5ece9ff81f870a35ae6658 Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Mon, 16 May 2016 09:35:13 -0700
Subject: [PATCH] Add context manager for Bigtable client.

This allows using the Bigtable Client class in a `with` statement.
---
 gcloud/bigtable/client.py      |  8 ++++++++
 gcloud/bigtable/test_client.py | 35 ++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/gcloud/bigtable/client.py b/gcloud/bigtable/client.py
index 073aec4e99c4..d04d6cf39510 100644
--- a/gcloud/bigtable/client.py
+++ b/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.
 
diff --git a/gcloud/bigtable/test_client.py b/gcloud/bigtable/test_client.py
index 1516297b5cba..1b13ef804cf9 100644
--- a/gcloud/bigtable/test_client.py
+++ b/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'