diff --git a/gcloud/datastore/_implicit_environ.py b/gcloud/datastore/_implicit_environ.py index 3a4307865515..ef71cec5a7de 100644 --- a/gcloud/datastore/_implicit_environ.py +++ b/gcloud/datastore/_implicit_environ.py @@ -86,6 +86,16 @@ def compute_engine_id(): connection.close() +def _get_production_dataset_id(): + """Gets the production application ID if it can be inferred.""" + return os.getenv(_DATASET_ENV_VAR_NAME) + + +def _get_gcd_dataset_id(): + """Gets the GCD application ID if it can be inferred.""" + return os.getenv(_GCD_DATASET_ENV_VAR_NAME) + + def _determine_default_dataset_id(dataset_id=None): """Determine default dataset ID explicitly or implicitly as fall-back. @@ -104,10 +114,10 @@ def _determine_default_dataset_id(dataset_id=None): :returns: Default dataset ID if it can be determined. """ if dataset_id is None: - dataset_id = os.getenv(_DATASET_ENV_VAR_NAME) + dataset_id = _get_production_dataset_id() if dataset_id is None: - dataset_id = os.getenv(_GCD_DATASET_ENV_VAR_NAME) + dataset_id = _get_gcd_dataset_id() if dataset_id is None: dataset_id = app_engine_id() diff --git a/gcloud/datastore/test__implicit_environ.py b/gcloud/datastore/test__implicit_environ.py index 8fab07047d7d..6d11ba2a7c87 100644 --- a/gcloud/datastore/test__implicit_environ.py +++ b/gcloud/datastore/test__implicit_environ.py @@ -65,260 +65,254 @@ def test_preset(self): self.assertEqual(self._callFUT(), SENTINEL) -class Test_set_default_dataset_id(unittest2.TestCase): +class Test__get_production_dataset_id(unittest2.TestCase): - def setUp(self): - from gcloud.datastore._testing import _setup_defaults - _setup_defaults(self) + def _callFUT(self): + from gcloud.datastore import _implicit_environ + return _implicit_environ._get_production_dataset_id() - def tearDown(self): - from gcloud.datastore._testing import _tear_down_defaults - _tear_down_defaults(self) + def test_no_value(self): + import os + from gcloud._testing import _Monkey - def _callFUT(self, dataset_id=None): - from gcloud.datastore._implicit_environ import set_default_dataset_id - return set_default_dataset_id(dataset_id=dataset_id) + environ = {} + with _Monkey(os, getenv=environ.get): + dataset_id = self._callFUT() + self.assertEqual(dataset_id, None) - def _monkeyEnviron(self, implicit_dataset_id, environ=None): + def test_value_set(self): import os from gcloud._testing import _Monkey from gcloud.datastore._implicit_environ import _DATASET_ENV_VAR_NAME - environ = environ or {_DATASET_ENV_VAR_NAME: implicit_dataset_id} - return _Monkey(os, getenv=environ.get) - def _monkeyImplicit(self, connection=None, app_identity=None): - from gcloud._testing import _Monkey - from gcloud.datastore import _implicit_environ - - if connection is None: - connection = _HTTPConnection(404, None) + MOCK_DATASET_ID = object() + environ = {_DATASET_ENV_VAR_NAME: MOCK_DATASET_ID} + with _Monkey(os, getenv=environ.get): + dataset_id = self._callFUT() + self.assertEqual(dataset_id, MOCK_DATASET_ID) - def _factory(host, timeout): - connection.host = host - connection.timeout = timeout - return connection - return _Monkey(_implicit_environ, - HTTPConnection=_factory, - app_identity=app_identity) +class Test__get_gcd_dataset_id(unittest2.TestCase): - def test_no_env_var_set(self): + def _callFUT(self): from gcloud.datastore import _implicit_environ + return _implicit_environ._get_gcd_dataset_id() - with self._monkeyEnviron(None): - with self._monkeyImplicit(): - self.assertRaises(EnvironmentError, self._callFUT) + def test_no_value(self): + import os + from gcloud._testing import _Monkey - self.assertEqual(_implicit_environ.get_default_dataset_id(), None) + environ = {} + with _Monkey(os, getenv=environ.get): + dataset_id = self._callFUT() + self.assertEqual(dataset_id, None) - def test_set_from_env_var(self): + def test_value_set(self): + import os + from gcloud._testing import _Monkey from gcloud.datastore import _implicit_environ - IMPLICIT_DATASET_ID = 'IMPLICIT' - - with self._monkeyEnviron(IMPLICIT_DATASET_ID): - with self._monkeyImplicit(): - self._callFUT() - self.assertEqual(_implicit_environ.get_default_dataset_id(), - IMPLICIT_DATASET_ID) - - def test_set_explicit_w_env_var_set(self): - from gcloud.datastore import _implicit_environ - EXPLICIT_DATASET_ID = 'EXPLICIT' + MOCK_DATASET_ID = object() + environ = { + _implicit_environ._GCD_DATASET_ENV_VAR_NAME: MOCK_DATASET_ID, + } + with _Monkey(os, getenv=environ.get): + dataset_id = self._callFUT() + self.assertEqual(dataset_id, MOCK_DATASET_ID) - with self._monkeyEnviron(None): - with self._monkeyImplicit(): - self._callFUT(EXPLICIT_DATASET_ID) - self.assertEqual(_implicit_environ.get_default_dataset_id(), - EXPLICIT_DATASET_ID) +class Test_app_engine_id(unittest2.TestCase): - def test_set_explicit_no_env_var_set(self): + def _callFUT(self): from gcloud.datastore import _implicit_environ - IMPLICIT_DATASET_ID = 'IMPLICIT' - EXPLICIT_DATASET_ID = 'EXPLICIT' + return _implicit_environ.app_engine_id() - with self._monkeyEnviron(IMPLICIT_DATASET_ID): - with self._monkeyImplicit(): - self._callFUT(EXPLICIT_DATASET_ID) + def test_no_value(self): + from gcloud._testing import _Monkey + from gcloud.datastore import _implicit_environ - self.assertEqual(_implicit_environ.get_default_dataset_id(), - EXPLICIT_DATASET_ID) + with _Monkey(_implicit_environ, app_identity=None): + dataset_id = self._callFUT() + self.assertEqual(dataset_id, None) - def test_set_explicit_None_wo_env_var_set(self): + def test_value_set(self): + from gcloud._testing import _Monkey from gcloud.datastore import _implicit_environ - with self._monkeyEnviron(None): - with self._monkeyImplicit(): - self.assertRaises(EnvironmentError, self._callFUT, None) - - self.assertEqual(_implicit_environ.get_default_dataset_id(), None) + APP_ENGINE_ID = object() + APP_IDENTITY = _AppIdentity(APP_ENGINE_ID) + with _Monkey(_implicit_environ, app_identity=APP_IDENTITY): + dataset_id = self._callFUT() + self.assertEqual(dataset_id, APP_ENGINE_ID) - def test_set_explicit_None_w_env_var_set(self): - from gcloud.datastore import _implicit_environ - IMPLICIT_DATASET_ID = 'IMPLICIT' - with self._monkeyEnviron(IMPLICIT_DATASET_ID): - with self._monkeyImplicit(): - self._callFUT(None) +class Test_compute_engine_id(unittest2.TestCase): - self.assertEqual(_implicit_environ.get_default_dataset_id(), - IMPLICIT_DATASET_ID) + def _callFUT(self): + from gcloud.datastore import _implicit_environ + return _implicit_environ.compute_engine_id() - def test_set_from_gcd_env_var(self): + def _monkeyConnection(self, connection): + from gcloud._testing import _Monkey from gcloud.datastore import _implicit_environ - GCD_ENV = _implicit_environ._GCD_DATASET_ENV_VAR_NAME - GCD_DATASET_ID = 'GCD-IMPLICIT' - ENVIRON = {GCD_ENV: GCD_DATASET_ID} + def _factory(host, timeout): + connection.host = host + connection.timeout = timeout + return connection - with self._monkeyEnviron(None, environ=ENVIRON): - with self._monkeyImplicit(): - self._callFUT() + return _Monkey(_implicit_environ, HTTPConnection=_factory) - self.assertEqual(_implicit_environ.get_default_dataset_id(), - GCD_DATASET_ID) + def test_bad_status(self): + connection = _HTTPConnection(404, None) + with self._monkeyConnection(connection): + dataset_id = self._callFUT() + self.assertEqual(dataset_id, None) - def test_set_gcd_and_production_env_vars(self): - from gcloud.datastore import _implicit_environ - from gcloud.datastore._implicit_environ import _DATASET_ENV_VAR_NAME + def test_success(self): + COMPUTE_ENGINE_ID = object() + connection = _HTTPConnection(200, COMPUTE_ENGINE_ID) + with self._monkeyConnection(connection): + dataset_id = self._callFUT() + self.assertEqual(dataset_id, COMPUTE_ENGINE_ID) - GCD_ENV = _implicit_environ._GCD_DATASET_ENV_VAR_NAME - IMPLICIT_DATASET_ID = 'IMPLICIT' - GCD_DATASET_ID = 'GCD-IMPLICIT' - ENVIRON = { - _DATASET_ENV_VAR_NAME: IMPLICIT_DATASET_ID, - GCD_ENV: GCD_DATASET_ID, - } + def test_socket_raises(self): + connection = _TimeoutHTTPConnection() + with self._monkeyConnection(connection): + dataset_id = self._callFUT() + self.assertEqual(dataset_id, None) - with self._monkeyEnviron(None, environ=ENVIRON): - with self._monkeyImplicit(): - self._callFUT() - self.assertEqual(_implicit_environ.get_default_dataset_id(), - IMPLICIT_DATASET_ID) +class Test__determine_default_dataset_id(unittest2.TestCase): - def test_set_gcd_env_vars_and_appengine(self): + def _callFUT(self, dataset_id=None): from gcloud.datastore import _implicit_environ + return _implicit_environ._determine_default_dataset_id( + dataset_id=dataset_id) - GCD_ENV = _implicit_environ._GCD_DATASET_ENV_VAR_NAME - GCD_DATASET_ID = 'GCD-IMPLICIT' - ENVIRON = {GCD_ENV: GCD_DATASET_ID} + def _determine_default_helper(self, prod=None, gcd=None, gae=None, + gce=None, dataset_id=None): + from gcloud._testing import _Monkey + from gcloud.datastore import _implicit_environ - APP_ENGINE_ID = 'GAE' - APP_IDENTITY = _AppIdentity(APP_ENGINE_ID) + _callers = [] - with self._monkeyEnviron(None, environ=ENVIRON): - with self._monkeyImplicit(app_identity=APP_IDENTITY): - self._callFUT() + def prod_mock(): + _callers.append('prod_mock') + return prod - self.assertEqual(_implicit_environ.get_default_dataset_id(), - GCD_DATASET_ID) + def gcd_mock(): + _callers.append('gcd_mock') + return gcd - def test_set_implicit_from_appengine(self): - from gcloud.datastore import _implicit_environ + def gae_mock(): + _callers.append('gae_mock') + return gae - APP_ENGINE_ID = 'GAE' - APP_IDENTITY = _AppIdentity(APP_ENGINE_ID) + def gce_mock(): + _callers.append('gce_mock') + return gce - with self._monkeyEnviron(None): - with self._monkeyImplicit(app_identity=APP_IDENTITY): - self._callFUT() + patched_methods = { + '_get_production_dataset_id': prod_mock, + '_get_gcd_dataset_id': gcd_mock, + 'app_engine_id': gae_mock, + 'compute_engine_id': gce_mock, + } - self.assertEqual(_implicit_environ.get_default_dataset_id(), - APP_ENGINE_ID) + with _Monkey(_implicit_environ, **patched_methods): + returned_dataset_id = self._callFUT(dataset_id) + + return returned_dataset_id, _callers + + def test_no_value(self): + dataset_id, callers = self._determine_default_helper() + self.assertEqual(dataset_id, None) + self.assertEqual(callers, + ['prod_mock', 'gcd_mock', 'gae_mock', 'gce_mock']) + + def test_explicit(self): + DATASET_ID = object() + dataset_id, callers = self._determine_default_helper( + dataset_id=DATASET_ID) + self.assertEqual(dataset_id, DATASET_ID) + self.assertEqual(callers, []) + + def test_prod(self): + DATASET_ID = object() + dataset_id, callers = self._determine_default_helper(prod=DATASET_ID) + self.assertEqual(dataset_id, DATASET_ID) + self.assertEqual(callers, ['prod_mock']) + + def test_gcd(self): + DATASET_ID = object() + dataset_id, callers = self._determine_default_helper(gcd=DATASET_ID) + self.assertEqual(dataset_id, DATASET_ID) + self.assertEqual(callers, ['prod_mock', 'gcd_mock']) + + def test_gae(self): + DATASET_ID = object() + dataset_id, callers = self._determine_default_helper(gae=DATASET_ID) + self.assertEqual(dataset_id, DATASET_ID) + self.assertEqual(callers, ['prod_mock', 'gcd_mock', 'gae_mock']) + + def test_gce(self): + DATASET_ID = object() + dataset_id, callers = self._determine_default_helper(gce=DATASET_ID) + self.assertEqual(dataset_id, DATASET_ID) + self.assertEqual(callers, + ['prod_mock', 'gcd_mock', 'gae_mock', 'gce_mock']) - def test_set_implicit_both_env_and_appengine(self): - from gcloud.datastore import _implicit_environ - IMPLICIT_DATASET_ID = 'IMPLICIT' - APP_IDENTITY = _AppIdentity('GAE') +class Test_set_default_dataset_id(unittest2.TestCase): + + def setUp(self): + from gcloud.datastore._testing import _setup_defaults + _setup_defaults(self) - with self._monkeyEnviron(IMPLICIT_DATASET_ID): - with self._monkeyImplicit(app_identity=APP_IDENTITY): - self._callFUT() + def tearDown(self): + from gcloud.datastore._testing import _tear_down_defaults + _tear_down_defaults(self) - self.assertEqual(_implicit_environ.get_default_dataset_id(), - IMPLICIT_DATASET_ID) + def _callFUT(self, dataset_id=None): + from gcloud.datastore._implicit_environ import set_default_dataset_id + return set_default_dataset_id(dataset_id=dataset_id) - def _implicit_compute_engine_helper(self, status): + def test_raises(self): + from gcloud._testing import _Monkey from gcloud.datastore import _implicit_environ - COMPUTE_ENGINE_ID = 'GCE' - if status == 200: - EXPECTED_ID = COMPUTE_ENGINE_ID - else: - EXPECTED_ID = None - - if status == 'RAISE': - connection = _TimeoutHTTPConnection() - else: - connection = _HTTPConnection(status, EXPECTED_ID) - - with self._monkeyEnviron(None): - with self._monkeyImplicit(connection=connection): - if EXPECTED_ID is None: - self.assertRaises(EnvironmentError, self._callFUT) - else: - self._callFUT() - - self.assertEqual(_implicit_environ.get_default_dataset_id(), - EXPECTED_ID) - self.assertEqual(connection.host, '169.254.169.254') - self.assertEqual(connection.timeout, 0.1) - self.assertEqual( - connection._called_args, - [('GET', '/computeMetadata/v1/project/project-id')]) - expected_kwargs = { - 'headers': { - 'Metadata-Flavor': 'Google', - }, - } - self.assertEqual(connection._called_kwargs, [expected_kwargs]) - self.assertEqual(connection._close_count, 1) + _called_dataset_id = [] - def test_set_implicit_from_compute_engine(self): - self._implicit_compute_engine_helper(200) + def mock_determine(dataset_id): + _called_dataset_id.append(dataset_id) + return None - def test_set_implicit_from_compute_engine_bad_status(self): - self._implicit_compute_engine_helper(404) + with _Monkey(_implicit_environ, + _determine_default_dataset_id=mock_determine): + self.assertRaises(EnvironmentError, self._callFUT) - def test_set_implicit_from_compute_engine_raise_timeout(self): - self._implicit_compute_engine_helper('RAISE') + self.assertEqual(_called_dataset_id, [None]) - def test_set_implicit_both_appengine_and_compute(self): + def test_set_correctly(self): + from gcloud._testing import _Monkey from gcloud.datastore import _implicit_environ - APP_ENGINE_ID = 'GAE' - APP_IDENTITY = _AppIdentity(APP_ENGINE_ID) - connection = _HTTPConnection(200, 'GCE') + self.assertEqual(_implicit_environ._DEFAULTS.dataset_id, None) - with self._monkeyEnviron(None): - with self._monkeyImplicit(connection=connection, - app_identity=APP_IDENTITY): - self._callFUT() + DATASET_ID = object() + _called_dataset_id = [] - self.assertEqual(_implicit_environ.get_default_dataset_id(), - APP_ENGINE_ID) - self.assertEqual(connection.host, None) - self.assertEqual(connection.timeout, None) + def mock_determine(dataset_id): + _called_dataset_id.append(dataset_id) + return DATASET_ID - def test_set_implicit_three_env_appengine_and_compute(self): - from gcloud.datastore import _implicit_environ - - IMPLICIT_DATASET_ID = 'IMPLICIT' - APP_IDENTITY = _AppIdentity('GAE') - connection = _HTTPConnection(200, 'GCE') - - with self._monkeyEnviron(IMPLICIT_DATASET_ID): - with self._monkeyImplicit(connection=connection, - app_identity=APP_IDENTITY): - self._callFUT() + with _Monkey(_implicit_environ, + _determine_default_dataset_id=mock_determine): + self._callFUT() - self.assertEqual(_implicit_environ.get_default_dataset_id(), - IMPLICIT_DATASET_ID) - self.assertEqual(connection.host, None) - self.assertEqual(connection.timeout, None) + self.assertEqual(_implicit_environ._DEFAULTS.dataset_id, DATASET_ID) + self.assertEqual(_called_dataset_id, [None]) class Test__lazy_property_deco(unittest2.TestCase):