Move '_require_connection' def into related '_implicit_environ' module.

gcloud/datastore/_implicit_environ.py

@@ -24,6 +24,7 @@
 from gcloud._helpers import _compute_engine_id
 from gcloud._helpers import _lazy_property_deco
 from gcloud.datastore.connection import Connection
+from gcloud.datastore.connection import _CONNECTIONS
 
 
 _DATASET_ENV_VAR_NAME = 'GCLOUD_DATASET_ID'
@@ -104,6 +105,28 @@ def get_default_dataset_id():
     return _DEFAULTS.dataset_id
 
 
+def _require_connection(connection=None):
+    """Infer a connection from the environment, if not passed explicitly.
+
+    :type connection: :class:`gcloud.datastore.connection.Connection`
+    :param connection: Optional.
+
+    :rtype: :class:`gcloud.datastore.connection.Connection`
+    :returns: A connection based on the current environment.
+    :raises: :class:`EnvironmentError` if ``connection`` is ``None``, and
+             cannot be inferred from the environment.
+    """
+    if connection is None:
+        top = _CONNECTIONS.top
+        if top is not None:
+            connection = top.connection
+        else:
+            connection = get_default_connection()
+            if connection is None:
+                raise EnvironmentError('Connection could not be inferred.')
+    return connection
+
+
 def set_default_connection(connection=None):
     """Set default connection either explicitly or implicitly as fall-back.
 

gcloud/datastore/_testing.py

@@ -31,3 +31,36 @@ def _setup_defaults(test_case, *args, **kwargs):
 
 def _tear_down_defaults(test_case):
     _implicit_environ._DEFAULTS = test_case._replaced_defaults
+
+
+class _NoCommitBatch(object):
+
+    def __init__(self, dataset_id, connection):
+        from gcloud.datastore.batch import Batch
+        self._batch = Batch(dataset_id, connection)
+
+    def __enter__(self):
+        from gcloud.datastore.connection import _CONNECTIONS
+        _CONNECTIONS.push(self._batch)
+        return self._batch
+
+    def __exit__(self, *args):
+        from gcloud.datastore.connection import _CONNECTIONS
+        _CONNECTIONS.pop()
+
+
+class _NoCommitTransaction(object):
+
+    def __init__(self, dataset_id, connection, transaction_id='TRANSACTION'):
+        from gcloud.datastore.transaction import Transaction
+        xact = self._transaction = Transaction(dataset_id, connection)
+        xact._id = transaction_id
+
+    def __enter__(self):
+        from gcloud.datastore.connection import _CONNECTIONS
+        _CONNECTIONS.push(self._transaction)
+        return self._transaction
+
+    def __exit__(self, *args):
+        from gcloud.datastore.connection import _CONNECTIONS
+        _CONNECTIONS.pop()

gcloud/datastore/api.py

@@ -66,28 +66,6 @@ def _require_dataset_id(dataset_id=None, first_key=None):
     return dataset_id
 
 
-def _require_connection(connection=None):
-    """Infer a connection from the environment, if not passed explicitly.
-
-    :type connection: :class:`gcloud.datastore.connection.Connection`
-    :param connection: Optional.
-
-    :rtype: :class:`gcloud.datastore.connection.Connection`
-    :returns: A connection based on the current environment.
-    :raises: :class:`EnvironmentError` if ``connection`` is ``None``, and
-             cannot be inferred from the environment.
-    """
-    if connection is None:
-        top = _CONNECTIONS.top
-        if top is not None:
-            connection = top.connection
-        else:
-            connection = _implicit_environ.get_default_connection()
-            if connection is None:
-                raise EnvironmentError('Connection could not be inferred.')
-    return connection
-
-
 def _extended_lookup(connection, dataset_id, key_pbs,
                      missing=None, deferred=None,
                      eventual=False, transaction_id=None):
@@ -201,7 +179,7 @@ def get(keys, missing=None, deferred=None, connection=None, dataset_id=None):
     if not keys:
         return []
 
-    connection = _require_connection(connection)
+    connection = _implicit_environ._require_connection(connection)
     dataset_id = _require_dataset_id(dataset_id, keys[0])
 
     if list(set([key.dataset_id for key in keys])) != [dataset_id]:
@@ -260,7 +238,7 @@ def put(entities, connection=None, dataset_id=None):
     if not entities:
         return
 
-    connection = _require_connection(connection)
+    connection = _implicit_environ._require_connection(connection)
     dataset_id = _require_dataset_id(dataset_id, entities[0].key)
 
     current = _CONNECTIONS.top
@@ -295,7 +273,7 @@ def delete(keys, connection=None, dataset_id=None):
     if not keys:
         return
 
-    connection = _require_connection(connection)
+    connection = _implicit_environ._require_connection(connection)
     dataset_id = _require_dataset_id(dataset_id, keys[0])
 
     # We allow partial keys to attempt a delete, the backend will fail.
@@ -325,7 +303,7 @@ def allocate_ids(incomplete_key, num_ids, connection=None):
     :returns: The (complete) keys allocated with ``incomplete_key`` as root.
     :raises: :class:`ValueError` if ``incomplete_key`` is not a partial key.
     """
-    connection = _require_connection(connection)
+    connection = _implicit_environ._require_connection(connection)
 
     if not incomplete_key.is_partial:
         raise ValueError(('Key is not partial.', incomplete_key))

gcloud/datastore/test__implicit_environ.py

@@ -299,6 +299,58 @@ def test_descriptor_for_connection(self):
             'connection' in _implicit_environ._DEFAULTS.__dict__)
 
 
+class Test__require_connection(unittest2.TestCase):
+
+    _MARKER = object()
+
+    def _callFUT(self, passed=_MARKER):
+        from gcloud.datastore._implicit_environ import _require_connection
+        if passed is self._MARKER:
+            return _require_connection()
+        return _require_connection(passed)
+
+    def _monkey(self, connection):
+        from gcloud.datastore._testing import _monkey_defaults
+        return _monkey_defaults(connection=connection)
+
+    def test_implicit_unset(self):
+        with self._monkey(None):
+            with self.assertRaises(EnvironmentError):
+                self._callFUT()
+
+    def test_implicit_unset_w_existing_batch(self):
+        from gcloud.datastore._testing import _NoCommitBatch
+        ID = 'DATASET'
+        CONNECTION = object()
+        with self._monkey(None):
+            with _NoCommitBatch(dataset_id=ID, connection=CONNECTION):
+                self.assertEqual(self._callFUT(), CONNECTION)
+
+    def test_implicit_unset_w_existing_transaction(self):
+        from gcloud.datastore._testing import _NoCommitTransaction
+        ID = 'DATASET'
+        CONNECTION = object()
+        with self._monkey(None):
+            with _NoCommitTransaction(dataset_id=ID, connection=CONNECTION):
+                self.assertEqual(self._callFUT(), CONNECTION)
+
+    def test_implicit_unset_passed_explicitly(self):
+        CONNECTION = object()
+        with self._monkey(None):
+            self.assertTrue(self._callFUT(CONNECTION) is CONNECTION)
+
+    def test_implicit_set(self):
+        IMPLICIT_CONNECTION = object()
+        with self._monkey(IMPLICIT_CONNECTION):
+            self.assertTrue(self._callFUT() is IMPLICIT_CONNECTION)
+
+    def test_implicit_set_passed_explicitly(self):
+        IMPLICIT_CONNECTION = object()
+        CONNECTION = object()
+        with self._monkey(IMPLICIT_CONNECTION):
+            self.assertTrue(self._callFUT(CONNECTION) is CONNECTION)
+
+
 class Test_set_default_connection(unittest2.TestCase):
 
     def setUp(self):

gcloud/datastore/test_api.py

@@ -41,27 +41,31 @@ def test_implicit_unset_w_keys(self):
             self.assertEqual(self._callFUT(first_key=_Key(ID)), ID)
 
     def test_implicit_unset_w_existing_batch_wo_keys(self):
+        from gcloud.datastore._testing import _NoCommitBatch
         ID = 'DATASET'
         with self._monkey(None):
             with _NoCommitBatch(dataset_id=ID, connection=object()):
                 self.assertEqual(self._callFUT(), ID)
 
     def test_implicit_unset_w_existing_batch_w_keys(self):
         from gcloud.datastore.test_batch import _Key
+        from gcloud.datastore._testing import _NoCommitBatch
         ID = 'DATASET'
         OTHER = 'OTHER'
         with self._monkey(None):
             with _NoCommitBatch(dataset_id=ID, connection=object()):
                 self.assertEqual(self._callFUT(first_key=_Key(OTHER)), ID)
 
     def test_implicit_unset_w_existing_transaction_wo_keys(self):
+        from gcloud.datastore._testing import _NoCommitTransaction
         ID = 'DATASET'
         with self._monkey(None):
             with _NoCommitTransaction(dataset_id=ID, connection=object()):
                 self.assertEqual(self._callFUT(), ID)
 
     def test_implicit_unset_w_existing_transaction_w_keys(self):
         from gcloud.datastore.test_batch import _Key
+        from gcloud.datastore._testing import _NoCommitTransaction
         ID = 'DATASET'
         OTHER = 'OTHER'
         with self._monkey(None):
@@ -108,56 +112,6 @@ def test_id_implicit_set_passed_explicitly_w_keys(self):
             self.assertEqual(self._callFUT(ID, first_key=_Key(OTHER)), ID)
 
 
-class Test__require_connection(unittest2.TestCase):
-
-    _MARKER = object()
-
-    def _callFUT(self, passed=_MARKER):
-        from gcloud.datastore.api import _require_connection
-        if passed is self._MARKER:
-            return _require_connection()
-        return _require_connection(passed)
-
-    def _monkey(self, connection):
-        from gcloud.datastore._testing import _monkey_defaults
-        return _monkey_defaults(connection=connection)
-
-    def test_implicit_unset(self):
-        with self._monkey(None):
-            with self.assertRaises(EnvironmentError):
-                self._callFUT()
-
-    def test_implicit_unset_w_existing_batch(self):
-        ID = 'DATASET'
-        CONNECTION = object()
-        with self._monkey(None):
-            with _NoCommitBatch(dataset_id=ID, connection=CONNECTION):
-                self.assertEqual(self._callFUT(), CONNECTION)
-
-    def test_implicit_unset_w_existing_transaction(self):
-        ID = 'DATASET'
-        CONNECTION = object()
-        with self._monkey(None):
-            with _NoCommitTransaction(dataset_id=ID, connection=CONNECTION):
-                self.assertEqual(self._callFUT(), CONNECTION)
-
-    def test_implicit_unset_passed_explicitly(self):
-        CONNECTION = object()
-        with self._monkey(None):
-            self.assertTrue(self._callFUT(CONNECTION) is CONNECTION)
-
-    def test_implicit_set(self):
-        IMPLICIT_CONNECTION = object()
-        with self._monkey(IMPLICIT_CONNECTION):
-            self.assertTrue(self._callFUT() is IMPLICIT_CONNECTION)
-
-    def test_implicit_set_passed_explicitly(self):
-        IMPLICIT_CONNECTION = object()
-        CONNECTION = object()
-        with self._monkey(IMPLICIT_CONNECTION):
-            self.assertTrue(self._callFUT(CONNECTION) is CONNECTION)
-
-
 class Test_get_function(unittest2.TestCase):
 
     def setUp(self):
@@ -499,6 +453,7 @@ def test_implicit_wo_transaction(self):
     def test_w_transaction(self):
         from gcloud.datastore.key import Key
         from gcloud.datastore.test_connection import _Connection
+        from gcloud.datastore._testing import _NoCommitTransaction
 
         DATASET_ID = 'DATASET'
         KIND = 'Kind'
@@ -661,6 +616,7 @@ def test_existing_batch_w_completed_key(self):
         from gcloud.datastore.test_batch import _Connection
         from gcloud.datastore.test_batch import _Entity
         from gcloud.datastore.test_batch import _Key
+        from gcloud.datastore._testing import _NoCommitBatch
 
         # Build basic mocks needed to delete.
         _DATASET = 'DATASET'
@@ -687,6 +643,7 @@ def test_implicit_connection(self):
         from gcloud.datastore.test_batch import _Connection
         from gcloud.datastore.test_batch import _Entity
         from gcloud.datastore.test_batch import _Key
+        from gcloud.datastore._testing import _NoCommitBatch
 
         # Build basic mocks needed to delete.
         _DATASET = 'DATASET'
@@ -804,6 +761,7 @@ def test_wo_batch_w_key_different_than_default_dataset_id(self):
     def test_w_existing_batch(self):
         from gcloud.datastore.test_batch import _Connection
         from gcloud.datastore.test_batch import _Key
+        from gcloud.datastore._testing import _NoCommitBatch
 
         # Build basic mocks needed to delete.
         _DATASET = 'DATASET'
@@ -825,6 +783,7 @@ def test_w_existing_batch(self):
     def test_w_existing_transaction(self):
         from gcloud.datastore.test_batch import _Connection
         from gcloud.datastore.test_batch import _Key
+        from gcloud.datastore._testing import _NoCommitTransaction
 
         # Build basic mocks needed to delete.
         _DATASET = 'DATASET'
@@ -847,6 +806,7 @@ def test_implicit_connection_and_dataset_id(self):
         from gcloud.datastore._testing import _monkey_defaults
         from gcloud.datastore.test_batch import _Connection
         from gcloud.datastore.test_batch import _Key
+        from gcloud.datastore._testing import _NoCommitBatch
 
         # Build basic mocks needed to delete.
         _DATASET = 'DATASET'
@@ -918,39 +878,6 @@ def test_with_already_completed_key(self):
                               COMPLETE_KEY, 2)
 
 
-class _NoCommitBatch(object):
-
-    def __init__(self, dataset_id, connection):
-        from gcloud.datastore.batch import Batch
-        self._batch = Batch(dataset_id, connection)
-
-    def __enter__(self):
-        from gcloud.datastore.connection import _CONNECTIONS
-        _CONNECTIONS.push(self._batch)
-        return self._batch
-
-    def __exit__(self, *args):
-        from gcloud.datastore.connection import _CONNECTIONS
-        _CONNECTIONS.pop()
-
-
-class _NoCommitTransaction(object):
-
-    def __init__(self, dataset_id, connection, transaction_id='TRANSACTION'):
-        from gcloud.datastore.transaction import Transaction
-        xact = self._transaction = Transaction(dataset_id, connection)
-        xact._id = transaction_id
-
-    def __enter__(self):
-        from gcloud.datastore.connection import _CONNECTIONS
-        _CONNECTIONS.push(self._transaction)
-        return self._transaction
-
-    def __exit__(self, *args):
-        from gcloud.datastore.connection import _CONNECTIONS
-        _CONNECTIONS.pop()
-
-
 class _HttpMultiple(object):
 
     def __init__(self, *responses):

gcloud/datastore/test_transaction.py

@@ -73,7 +73,7 @@ def test_ctor_with_env(self):
         self.assertEqual(xact._status, self._getTargetClass()._INITIAL)
 
     def test_current(self):
-        from gcloud.datastore.test_api import _NoCommitBatch
+        from gcloud.datastore._testing import _NoCommitBatch
         _DATASET = 'DATASET'
         connection = _Connection()
         xact1 = self._makeOne(_DATASET, connection)