googleapis · dhermes · Mar 30, 2015 · Mar 30, 2015
diff --git a/gcloud/storage/__init__.py b/gcloud/storage/__init__.py
@@ -51,6 +51,7 @@
 from gcloud.storage.api import get_all_buckets
 from gcloud.storage.api import get_bucket
 from gcloud.storage.api import lookup_bucket
+from gcloud.storage.batch import Batch
 from gcloud.storage.blob import Blob
 from gcloud.storage.bucket import Bucket
 from gcloud.storage.connection import Connection

diff --git a/gcloud/storage/api.py b/gcloud/storage/api.py
@@ -21,6 +21,7 @@
 from gcloud.exceptions import NotFound
 from gcloud._helpers import get_default_project
 from gcloud.storage._implicit_environ import get_default_connection
+from gcloud.storage.batch import Batch
 from gcloud.storage.bucket import Bucket
 from gcloud.storage.iterator import Iterator
 
@@ -51,9 +52,7 @@ def lookup_bucket(bucket_name, connection=None):
     :rtype: :class:`gcloud.storage.bucket.Bucket`
     :returns: The bucket matching the name provided or None if not found.
     """
-    if connection is None:
-        connection = get_default_connection()
-
+    connection = _require_connection(connection)
     try:
         return get_bucket(bucket_name, connection=connection)
     except NotFound:
@@ -84,8 +83,7 @@ def get_all_buckets(project=None, connection=None):
     :rtype: iterable of :class:`gcloud.storage.bucket.Bucket` objects.
     :returns: All buckets belonging to this project.
     """
-    if connection is None:
-        connection = get_default_connection()
+    connection = _require_connection(connection)
     if project is None:
         project = get_default_project()
     extra_params = {'project': project}
@@ -122,9 +120,7 @@ def get_bucket(bucket_name, connection=None):
     :returns: The bucket matching the name provided.
     :raises: :class:`gcloud.exceptions.NotFound`
     """
-    if connection is None:
-        connection = get_default_connection()
-
+    connection = _require_connection(connection)
     bucket = Bucket(bucket_name, connection=connection)
     bucket.reload()
     return bucket
@@ -160,8 +156,7 @@ def create_bucket(bucket_name, project=None, connection=None):
     :raises: :class:`gcloud.exceptions.Conflict` if
              there is a confict (bucket already exists, invalid name, etc.)
     """
-    if connection is None:
-        connection = get_default_connection()
+    connection = _require_connection(connection)
     if project is None:
         project = get_default_project()
 
@@ -201,3 +196,27 @@ def get_items_from_response(self, response):
             bucket = Bucket(name, connection=self.connection)
             bucket._properties = item
             yield bucket
+
+
+def _require_connection(connection=None):
+    """Infer a connection from the environment, if not passed explicitly.
+
+    :type connection: :class:`gcloud.storage.connection.Connection`
+    :param connection: Optional.
+
+    :rtype: :class:`gcloud.storage.connection.Connection`
+    :returns: A connection based on the current environment.
+    :raises: :class:`EnvironmentError` if ``connection`` is ``None``, and
+             cannot be inferred from the environment.
+    """
+    # NOTE: We use current Batch directly since it inherits from Connection.
+    if connection is None:
+        connection = Batch.current()
+
+    if connection is None:
+        connection = get_default_connection()
+
+    if connection is None:
+        raise EnvironmentError('Connection could not be inferred.')
+
+    return connection
diff --git a/gcloud/storage/batch.py b/gcloud/storage/batch.py
@@ -155,6 +155,11 @@ def finish(self):
         self._responses = list(_unpack_batch_response(response, content))
         return self._responses
 
+    @staticmethod
+    def current():
+        """Return the topmost batch, or None."""
+        return _BATCHES.top
+
     def __enter__(self):
         _BATCHES.push(self)
         return self

diff --git a/gcloud/storage/test_api.py b/gcloud/storage/test_api.py
@@ -284,6 +284,44 @@ def test_get_items_from_response_non_empty(self):
         self.assertEqual(bucket.name, BLOB_NAME)
 
 
+class Test__require_connection(unittest2.TestCase):
+
+    def _callFUT(self, connection=None):
+        from gcloud.storage.api import _require_connection
+        return _require_connection(connection=connection)
+
+    def _monkey(self, connection):
+        from gcloud.storage._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):
+        CONNECTION = object()
+        with self._monkey(None):
+            with _NoCommitBatch(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 Http(object):
 
     _called_with = None
@@ -296,3 +334,18 @@ def __init__(self, headers, content):
     def request(self, **kw):
         self._called_with = kw
         return self._response, self._content
+
+
+class _NoCommitBatch(object):
+
+    def __init__(self, connection):
+        self._connection = connection
+
+    def __enter__(self):
+        from gcloud.storage.batch import _BATCHES
+        _BATCHES.push(self._connection)
+        return self._connection
+
+    def __exit__(self, *args):
+        from gcloud.storage.batch import _BATCHES
+        _BATCHES.pop()
diff --git a/regression/storage.py b/regression/storage.py
@@ -22,7 +22,6 @@
 from gcloud import storage
 from gcloud import _helpers
 from gcloud.storage._helpers import _base64_md5hash
-from gcloud.storage.batch import Batch
 
 
 HTTP = httplib2.Http()
@@ -52,7 +51,7 @@ def setUp(self):
         self.case_buckets_to_delete = []
 
     def tearDown(self):
-        with Batch() as batch:
+        with storage.Batch() as batch:
             for bucket_name in self.case_buckets_to_delete:
                 storage.Bucket(bucket_name, connection=batch).delete()