Merge pull request #478 from dhermes/make-default-datastore-connection

Address first part of 477: Implements datastore.set_default_connection.

Author: tseaver

gcloud/datastore/__init__.py

@@ -59,12 +59,10 @@
 
 
 def set_default_dataset(dataset_id=None):
-    """Determines auth settings from local enviroment.
+    """Set default dataset ID either explicitly or implicitly as fall-back.
 
-    Sets a default dataset either explicitly or via the local
-    enviroment. Currently only supports enviroment variable but will
-    implicitly support App Engine, Compute Engine and other environments
-    in the future.
+    In implicit case, currently only supports enviroment variable but will
+    support App Engine, Compute Engine and other environments in the future.
 
     Local environment variable used is:
     - GCLOUD_DATASET_ID
@@ -80,6 +78,16 @@ def set_default_dataset(dataset_id=None):
         _implicit_environ.DATASET = get_dataset(dataset_id)
 
 
+def set_default_connection(connection=None):
+    """Set default connection either explicitly or implicitly as fall-back.
+
+    :type connection: :class:`gcloud.datastore.connection.Connection`
+    :param connection: A connection provided to be the default.
+    """
+    connection = connection or get_connection()
+    _implicit_environ.CONNECTION = connection
+
+
 def get_connection():
     """Shortcut method to establish a connection to the Cloud Datastore.
 

gcloud/datastore/_implicit_environ.py

@@ -1,7 +1,7 @@
 """Module to provide implicit behavior based on enviroment.
 
 Acts as a mutable namespace to allow the datastore package to
-imply the current dataset from the enviroment.
+imply the current dataset and connection from the enviroment.
 
 Also provides a base class for classes in the `datastore` package
 which could utilize the implicit enviroment.
@@ -11,6 +11,9 @@
 DATASET = None
 """Module global to allow persistent implied dataset from enviroment."""
 
+CONNECTION = None
+"""Module global to allow persistent implied connection from enviroment."""
+
 
 class _DatastoreBase(object):
     """Base for all classes in the datastore package.

gcloud/datastore/test___init__.py

@@ -89,6 +89,43 @@ def test_set_explicit(self):
         self._test_with_environ({}, DATASET_ID, dataset_id=DATASET_ID)
 
 
+class Test_set_default_connection(unittest2.TestCase):
+
+    def setUp(self):
+        from gcloud.datastore import _implicit_environ
+        self._replaced_connection = _implicit_environ.CONNECTION
+        _implicit_environ.CONNECTION = None
+
+    def tearDown(self):
+        from gcloud.datastore import _implicit_environ
+        _implicit_environ.CONNECTION = self._replaced_connection
+
+    def _callFUT(self, connection=None):
+        from gcloud.datastore import set_default_connection
+        return set_default_connection(connection=connection)
+
+    def test_set_explicit(self):
+        from gcloud.datastore import _implicit_environ
+
+        self.assertEqual(_implicit_environ.CONNECTION, None)
+        fake_cnxn = object()
+        self._callFUT(connection=fake_cnxn)
+        self.assertEqual(_implicit_environ.CONNECTION, fake_cnxn)
+
+    def test_set_implicit(self):
+        from gcloud._testing import _Monkey
+        from gcloud import datastore
+        from gcloud.datastore import _implicit_environ
+
+        self.assertEqual(_implicit_environ.CONNECTION, None)
+
+        fake_cnxn = object()
+        with _Monkey(datastore, get_connection=lambda: fake_cnxn):
+            self._callFUT()
+
+        self.assertEqual(_implicit_environ.CONNECTION, fake_cnxn)
+
+
 class Test_get_dataset(unittest2.TestCase):
 
     def _callFUT(self, dataset_id):