Merge pull request googleapis#1474 from dhermes/bigtable-connection-open

dhermes · dhermes · commit 865125236d22 · 2016-02-15T17:46:25.000-08:00
Implementing open()/close() on HappyBase connection.
diff --git a/gcloud/bigtable/happybase/connection.py b/gcloud/bigtable/happybase/connection.py
@@ -127,6 +127,8 @@ class Connection(object):
              parameters are specified with a value other than the defaults.
     """
 
+    _cluster = None
+
     def __init__(self, timeout=None, autoconnect=True, table_prefix=None,
                  table_prefix_separator='_', cluster=None, **kwargs):
         self._handle_legacy_args(kwargs)
@@ -140,7 +142,6 @@ def __init__(self, timeout=None, autoconnect=True, table_prefix=None,
                             'received', table_prefix_separator,
                             type(table_prefix_separator))
 
-        self.autoconnect = autoconnect
         self.table_prefix = table_prefix
         self.table_prefix_separator = table_prefix_separator
 
@@ -152,6 +153,11 @@ def __init__(self, timeout=None, autoconnect=True, table_prefix=None,
                                  'cluster is passed')
             self._cluster = cluster.copy()
 
+        if autoconnect:
+            self.open()
+
+        self._initialized = True
+
     @staticmethod
     def _handle_legacy_args(arguments_dict):
         """Check legacy HappyBase arguments and warn if set.
@@ -174,3 +180,25 @@ def _handle_legacy_args(arguments_dict):
         if arguments_dict:
             unexpected_names = arguments_dict.keys()
             raise TypeError('Received unexpected arguments', unexpected_names)
+
+    def open(self):
+        """Open the underlying transport to Cloud Bigtable.
+
+        This method opens the underlying HTTP/2 gRPC connection using a
+        :class:`.Client` bound to the :class:`.Cluster` owned by
+        this connection.
+        """
+        self._cluster._client.start()
+
+    def close(self):
+        """Close the underlying transport to Cloud Bigtable.
+
+        This method closes the underlying HTTP/2 gRPC connection using a
+        :class:`.Client` bound to the :class:`.Cluster` owned by
+        this connection.
+        """
+        self._cluster._client.stop()
+
+    def __del__(self):
+        if self._cluster is not None:
+            self.close()
diff --git a/gcloud/bigtable/happybase/test_connection.py b/gcloud/bigtable/happybase/test_connection.py
@@ -81,13 +81,23 @@ def _makeOne(self, *args, **kwargs):
 
     def test_constructor_defaults(self):
         cluster = _Cluster()  # Avoid implicit environ check.
+        self.assertEqual(cluster._client.start_calls, 0)
         connection = self._makeOne(cluster=cluster)
+        self.assertEqual(cluster._client.start_calls, 1)
+        self.assertEqual(cluster._client.stop_calls, 0)
 
-        self.assertTrue(connection.autoconnect)
         self.assertEqual(connection._cluster, cluster)
         self.assertEqual(connection.table_prefix, None)
         self.assertEqual(connection.table_prefix_separator, '_')
 
+    def test_constructor_no_autoconnect(self):
+        cluster = _Cluster()  # Avoid implicit environ check.
+        connection = self._makeOne(autoconnect=False, cluster=cluster)
+        self.assertEqual(cluster._client.start_calls, 0)
+        self.assertEqual(cluster._client.stop_calls, 0)
+        self.assertEqual(connection.table_prefix, None)
+        self.assertEqual(connection.table_prefix_separator, '_')
+
     def test_constructor_missing_cluster(self):
         from gcloud._testing import _Monkey
         from gcloud.bigtable.happybase import connection as MUT
@@ -121,7 +131,6 @@ def test_constructor_explicit(self):
             table_prefix=table_prefix,
             table_prefix_separator=table_prefix_separator,
             cluster=cluster)
-        self.assertFalse(connection.autoconnect)
         self.assertEqual(connection.table_prefix, table_prefix)
         self.assertEqual(connection.table_prefix_separator,
                          table_prefix_separator)
@@ -172,6 +181,37 @@ def test_constructor_non_string_prefix_separator(self):
             self._makeOne(autoconnect=False,
                           table_prefix_separator=table_prefix_separator)
 
+    def test_open(self):
+        cluster = _Cluster()  # Avoid implicit environ check.
+        connection = self._makeOne(autoconnect=False, cluster=cluster)
+        self.assertEqual(cluster._client.start_calls, 0)
+        connection.open()
+        self.assertEqual(cluster._client.start_calls, 1)
+        self.assertEqual(cluster._client.stop_calls, 0)
+
+    def test_close(self):
+        cluster = _Cluster()  # Avoid implicit environ check.
+        connection = self._makeOne(autoconnect=False, cluster=cluster)
+        self.assertEqual(cluster._client.stop_calls, 0)
+        connection.close()
+        self.assertEqual(cluster._client.stop_calls, 1)
+        self.assertEqual(cluster._client.start_calls, 0)
+
+    def test___del__with_cluster(self):
+        cluster = _Cluster()  # Avoid implicit environ check.
+        connection = self._makeOne(autoconnect=False, cluster=cluster)
+        self.assertEqual(cluster._client.stop_calls, 0)
+        connection.__del__()
+        self.assertEqual(cluster._client.stop_calls, 1)
+
+    def test___del__no_cluster(self):
+        cluster = _Cluster()  # Avoid implicit environ check.
+        connection = self._makeOne(autoconnect=False, cluster=cluster)
+        self.assertEqual(cluster._client.stop_calls, 0)
+        del connection._cluster
+        connection.__del__()
+        self.assertEqual(cluster._client.stop_calls, 0)
+
 
 class _Client(object):
 
