From fa0412762b7d651c8e9366dd841302edf5ee58b3 Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Wed, 18 Feb 2015 19:08:45 -0800
Subject: [PATCH]  Adding lazy loading support for datastore connection.

---
 gcloud/datastore/_implicit_environ.py      | 17 +++++++++--
 gcloud/datastore/test__implicit_environ.py | 34 ++++++++++++----------
 2 files changed, 33 insertions(+), 18 deletions(-)

diff --git a/gcloud/datastore/_implicit_environ.py b/gcloud/datastore/_implicit_environ.py
index 30a6f487c4457..58a0a870cb653 100644
--- a/gcloud/datastore/_implicit_environ.py
+++ b/gcloud/datastore/_implicit_environ.py
@@ -119,9 +119,9 @@ def _determine_default_dataset_id(dataset_id=None):
 
 
 def _lazy_dataset_id():
-    """Alias wrapper for _determine_default_dataset_id.
+    """Alias wrapper for ``_determine_default_dataset_id``.
 
-    Unit test need to be able to replace _determine_default_dataset_id()
+    Unit tests need to be able to replace :func:`_determine_default_dataset_id`
     so we can't wrap the actual ``function`` object in a ``_LazyProperty``.
     """
     return _determine_default_dataset_id()
@@ -181,6 +181,15 @@ def get_connection():
     return Connection(credentials=scoped_credentials)
 
 
+def _lazy_connection():
+    """Alias wrapper for ``get_connection``.
+
+    Unit tests need to be able to replace :func:`get_connection`
+    so we can't wrap the actual ``function`` object in a ``_LazyProperty``.
+    """
+    return get_connection()
+
+
 def set_default_connection(connection=None):
     """Set default connection either explicitly or implicitly as fall-back.
 
@@ -236,9 +245,11 @@ class _DefaultsContainer(object):
     """
 
     dataset_id = _LazyProperty('dataset_id', _lazy_dataset_id)
+    connection = _LazyProperty('connection', _lazy_connection)
 
     def __init__(self, connection=None, dataset_id=None, implicit=False):
-        self.connection = connection
+        if connection is not None or not implicit:
+            self.connection = connection
         if dataset_id is not None or not implicit:
             self.dataset_id = dataset_id
 
diff --git a/gcloud/datastore/test__implicit_environ.py b/gcloud/datastore/test__implicit_environ.py
index 9e64c3d27861f..497d296f9c87b 100644
--- a/gcloud/datastore/test__implicit_environ.py
+++ b/gcloud/datastore/test__implicit_environ.py
@@ -325,7 +325,7 @@ def test_set_implicit_three_env_appengine_and_compute(self):
         self.assertEqual(connection.timeout, None)
 
 
-class Test_lazy_loaded_dataset_id(unittest2.TestCase):
+class Test_lazy_loading(unittest2.TestCase):
 
     def setUp(self):
         from gcloud.datastore._testing import _setup_defaults
@@ -335,15 +335,6 @@ def tearDown(self):
         from gcloud.datastore._testing import _tear_down_defaults
         _tear_down_defaults(self)
 
-    def test_prop_default(self):
-        from gcloud.datastore import _implicit_environ
-        from gcloud.datastore._implicit_environ import _DefaultsContainer
-        from gcloud.datastore._implicit_environ import _LazyProperty
-
-        self.assertTrue(isinstance(_DefaultsContainer.dataset_id,
-                                   _LazyProperty))
-        self.assertEqual(_implicit_environ._DEFAULTS.dataset_id, None)
-
     def test_prop_on_wrong_class(self):
         from gcloud.datastore._implicit_environ import _LazyProperty
 
@@ -357,7 +348,7 @@ class FakeEnv(object):
         self.assertTrue(FakeEnv.dataset_id is data_prop)
         self.assertTrue(FakeEnv().dataset_id is data_prop)
 
-    def test_prop_descriptor(self):
+    def test_descriptor_for_dataset_id(self):
         from gcloud._testing import _Monkey
         from gcloud.datastore import _implicit_environ
 
@@ -366,17 +357,30 @@ def test_prop_descriptor(self):
 
         DEFAULT = object()
 
-        def mock_default():
-            return DEFAULT
-
         with _Monkey(_implicit_environ,
-                     _determine_default_dataset_id=mock_default):
+                     _determine_default_dataset_id=lambda: DEFAULT):
             lazy_loaded = _implicit_environ._DEFAULTS.dataset_id
 
         self.assertEqual(lazy_loaded, DEFAULT)
         self.assertTrue(
             'dataset_id' in _implicit_environ._DEFAULTS.__dict__)
 
+    def test_descriptor_for_connection(self):
+        from gcloud._testing import _Monkey
+        from gcloud.datastore import _implicit_environ
+
+        self.assertFalse(
+            'connection' in _implicit_environ._DEFAULTS.__dict__)
+
+        DEFAULT = object()
+
+        with _Monkey(_implicit_environ, get_connection=lambda: DEFAULT):
+            lazy_loaded = _implicit_environ._DEFAULTS.connection
+
+        self.assertEqual(lazy_loaded, DEFAULT)
+        self.assertTrue(
+            'connection' in _implicit_environ._DEFAULTS.__dict__)
+
 
 class Test_get_connection(unittest2.TestCase):