Skip to content

Commit

Permalink
Add context manager for Bigtable client.
Browse files Browse the repository at this point in the history
This allows using the Bigtable Client class in a `with` statement.
  • Loading branch information
tswast committed May 16, 2016
1 parent 755bf52 commit 87aed7f
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 0 deletions.
8 changes: 8 additions & 0 deletions gcloud/bigtable/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand All @@ -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.
Expand Down
35 changes: 35 additions & 0 deletions gcloud/bigtable/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down

0 comments on commit 87aed7f

Please sign in to comment.