Adding support for GCD dev server in datastore package.

Enabling - set_default_dataset_id() support for DATASTORE_DATASET environment variable name - gcloud.datastore.Connection.API_BASE_URL override for DATASTORE_HOST environment variable name Fixes #650.
dhermes · Feb 18, 2015 · Feb 17, 2015 · Feb 18, 2015 · Feb 17, 2015 · d2aaab844cb2f4fb527eb108dc23edc9f5f3285c
commit d2aaab844cb2f4fb527eb108dc23edc9f5f3285c
diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py
@@ -67,6 +67,7 @@
 """The scopes required for authenticating as a Cloud Datastore consumer."""
 
 _DATASET_ENV_VAR_NAME = 'GCLOUD_DATASET_ID'
+_GCD_DATASET_ENV_VAR_NAME = 'DATASTORE_DATASET'
 
 
 def set_default_dataset_id(dataset_id=None):
@@ -86,6 +87,9 @@ def set_default_dataset_id(dataset_id=None):
     if dataset_id is None:
         dataset_id = os.getenv(_DATASET_ENV_VAR_NAME)
 
+    if dataset_id is None:
+        dataset_id = os.getenv(_GCD_DATASET_ENV_VAR_NAME)
+
     if dataset_id is None:
         dataset_id = _implicit_environ.app_engine_id()
 

diff --git a/gcloud/datastore/connection.py b/gcloud/datastore/connection.py
@@ -14,12 +14,17 @@
 
 """Connections to gcloud datastore API servers."""
 
+import os
+
 from gcloud import connection
 from gcloud.exceptions import make_exception
 from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
 from gcloud.datastore import helpers
 
 
+_GCD_HOST_ENV_VAR_NAME = 'DATASTORE_HOST'
+
+
 class Connection(connection.Connection):
     """A connection to the Google Cloud Datastore via the Protobuf API.
 
@@ -30,6 +35,10 @@ class Connection(connection.Connection):
     :param credentials: The OAuth2 Credentials to use for this connection.
     """
 
+    API_BASE_URL = os.getenv(_GCD_HOST_ENV_VAR_NAME,
+                             connection.Connection.API_BASE_URL)
+    """The base of the API call URL."""
+
     API_VERSION = 'v1beta2'
     """The version of the API, used in building the API call's URL."""
 

diff --git a/gcloud/datastore/test___init__.py b/gcloud/datastore/test___init__.py
@@ -30,11 +30,11 @@ def _callFUT(self, dataset_id=None):
         from gcloud.datastore import set_default_dataset_id
         return set_default_dataset_id(dataset_id=dataset_id)
 
-    def _monkeyEnviron(self, implicit_dataset_id):
+    def _monkeyEnviron(self, implicit_dataset_id, environ=None):
         import os
         from gcloud._testing import _Monkey
         from gcloud.datastore import _DATASET_ENV_VAR_NAME
-        environ = {_DATASET_ENV_VAR_NAME: implicit_dataset_id}
+        environ = environ or {_DATASET_ENV_VAR_NAME: implicit_dataset_id}
         return _Monkey(os, getenv=environ.get)
 
     def _monkeyImplicit(self, connection=None, app_identity=None):
@@ -112,6 +112,55 @@ def test_set_explicit_None_w_env_var_set(self):
 
         self.assertEqual(_implicit_environ.DATASET_ID, IMPLICIT_DATASET_ID)
 
+    def test_set_from_gcd_env_var(self):
+        from gcloud.datastore import _GCD_DATASET_ENV_VAR_NAME
+        from gcloud.datastore import _implicit_environ
+
+        GCD_DATASET_ID = 'GCD-IMPLICIT'
+        ENVIRON = {_GCD_DATASET_ENV_VAR_NAME: GCD_DATASET_ID}
+
+        with self._monkeyEnviron(None, environ=ENVIRON):
+            with self._monkeyImplicit():
+                self._callFUT()
+
+        self.assertEqual(_implicit_environ.DATASET_ID, GCD_DATASET_ID)
+
+    def test_set_gcd_and_production_env_vars(self):
+        from gcloud.datastore import _DATASET_ENV_VAR_NAME
+        from gcloud.datastore import _GCD_DATASET_ENV_VAR_NAME
+        from gcloud.datastore import _implicit_environ
+
+        IMPLICIT_DATASET_ID = 'IMPLICIT'
+        GCD_DATASET_ID = 'GCD-IMPLICIT'
+        ENVIRON = {
+            _DATASET_ENV_VAR_NAME: IMPLICIT_DATASET_ID,
+            _GCD_DATASET_ENV_VAR_NAME: GCD_DATASET_ID,
+        }
+
+        with self._monkeyEnviron(None, environ=ENVIRON):
+            with self._monkeyImplicit():
+                self._callFUT()
+
+        self.assertNotEqual(_implicit_environ.DATASET_ID, GCD_DATASET_ID)
+        self.assertEqual(_implicit_environ.DATASET_ID, IMPLICIT_DATASET_ID)
+
+    def test_set_gcd_env_vars_and_appengine(self):
+        from gcloud.datastore import _GCD_DATASET_ENV_VAR_NAME
+        from gcloud.datastore import _implicit_environ
+
+        GCD_DATASET_ID = 'GCD-IMPLICIT'
+        ENVIRON = {_GCD_DATASET_ENV_VAR_NAME: GCD_DATASET_ID}
+
+        APP_ENGINE_ID = 'GAE'
+        APP_IDENTITY = _AppIdentity(APP_ENGINE_ID)
+
+        with self._monkeyEnviron(None, environ=ENVIRON):
+            with self._monkeyImplicit(app_identity=APP_IDENTITY):
+                self._callFUT()
+
+        self.assertNotEqual(_implicit_environ.DATASET_ID, APP_ENGINE_ID)
+        self.assertEqual(_implicit_environ.DATASET_ID, GCD_DATASET_ID)
+
     def test_set_implicit_from_appengine(self):
         from gcloud.datastore import _implicit_environ
 

diff --git a/gcloud/datastore/test_connection.py b/gcloud/datastore/test_connection.py
@@ -46,6 +46,47 @@ def _verifyProtobufCall(self, called_with, URI, conn):
         self.assertEqual(called_with['headers']['User-Agent'],
                          conn.USER_AGENT)
 
+    def test_default_url(self):
+        from gcloud.connection import Connection
+
+        klass = self._getTargetClass()
+        self.assertEqual(klass.API_BASE_URL, Connection.API_BASE_URL)
+
+    def test_custom_url(self):
+        import os
+        import sys
+        from gcloud._testing import _Monkey
+        from gcloud.connection import Connection as BaseConnection
+        from gcloud.datastore.connection import _GCD_HOST_ENV_VAR_NAME
+
+        HOST = object()
+        fake_environ = {_GCD_HOST_ENV_VAR_NAME: HOST}
+
+        def fake_getenv(key, default_val=None):
+            return fake_environ.get(key, default_val)
+
+        # We want to temporarily all the gcloud.datastore modules except for
+        # the gcloud.datastore package itself. Then re-import.
+        gcloud_keys = [key for key in sys.modules
+                       if 'gcloud.datastore.' in key]
+        modules_as_is = sys.modules.copy()
+        try:
+            for key in gcloud_keys:
+                sys.modules.pop(key)
+
+            with _Monkey(os, getenv=fake_getenv):
+                from gcloud.datastore.connection import Connection
+        finally:
+            sys.modules.update(modules_as_is)
+
+        self.assertNotEqual(Connection.API_BASE_URL,
+                            BaseConnection.API_BASE_URL)
+        self.assertEqual(Connection.API_BASE_URL, HOST)
+
+        # Make sure the restored import.
+        from gcloud.datastore.connection import Connection
+        self.assertEqual(Connection.API_BASE_URL, BaseConnection.API_BASE_URL)
+
     def test_ctor_defaults(self):
         conn = self._makeOne()
         self.assertEqual(conn.credentials, None)