Adding namespace to dataset class.

gcloud/datastore/dataset.py

@@ -31,14 +31,18 @@ class Dataset(object):
     :type dataset_id: string
     :param dataset_id: (required) dataset ID to pass to proxied API methods.
 
+    :type namespace: string
+    :param namespace: (optional) namespace to pass to proxied API methods.
+
     :type connection: :class:`gcloud.datastore.connection.Connection`, or None
     :param connection: (optional) connection to pass to proxied API methods
     """
 
-    def __init__(self, dataset_id, connection=None):
+    def __init__(self, dataset_id, namespace=None, connection=None):
         if dataset_id is None:
             raise ValueError('dataset_id required')
         self.dataset_id = dataset_id
+        self.namespace = namespace
         self.connection = connection
 
     def get(self, key, missing=None, deferred=None):
@@ -98,6 +102,8 @@ def key(self, *path_args, **kwargs):
         if 'dataset_id' in kwargs:
             raise TypeError('Cannot pass dataset_id')
         kwargs['dataset_id'] = self.dataset_id
+        if 'namespace' not in kwargs:
+            kwargs['namespace'] = self.namespace
         return Key(*path_args, **kwargs)
 
     def batch(self):
@@ -123,4 +129,6 @@ def query(self, **kwargs):
         if 'dataset_id' in kwargs:
             raise TypeError('Cannot pass dataset_id')
         kwargs['dataset_id'] = self.dataset_id
+        if 'namespace' not in kwargs:
+            kwargs['namespace'] = self.namespace
         return Query(**kwargs)

gcloud/datastore/test_dataset.py

@@ -23,8 +23,9 @@ def _getTargetClass(self):
         from gcloud.datastore.dataset import Dataset
         return Dataset
 
-    def _makeOne(self, dataset_id=DATASET_ID, connection=None):
-        return self._getTargetClass()(dataset_id, connection)
+    def _makeOne(self, dataset_id=DATASET_ID, namespace=None, connection=None):
+        return self._getTargetClass()(dataset_id, namespace=namespace,
+                                      connection=connection)
 
     def test_ctor_w_dataset_id_None(self):
         self.assertRaises(ValueError, self._makeOne, None)
@@ -33,11 +34,13 @@ def test_ctor_w_dataset_id_no_connection(self):
         dataset = self._makeOne()
         self.assertEqual(dataset.dataset_id, self.DATASET_ID)
 
-    def test_ctor_w_dataset_id_w_connection(self):
+    def test_ctor_w_explicit_inputs(self):
         conn = object()
-        dataset = self._makeOne(connection=conn)
+        namespace = object()
+        dataset = self._makeOne(namespace=namespace, connection=conn)
         self.assertEqual(dataset.dataset_id, self.DATASET_ID)
         self.assertTrue(dataset.connection is conn)
+        self.assertTrue(dataset.namespace is namespace)
 
     def test_get_defaults(self):
         from gcloud.datastore import dataset as MUT
@@ -294,7 +297,48 @@ def test_key_wo_dataset_id(self):
 
         self.assertTrue(isinstance(key, _Dummy))
         self.assertEqual(key.args, (KIND, ID))
-        self.assertEqual(key.kwargs, {'dataset_id': self.DATASET_ID})
+        expected_kwargs = {
+            'dataset_id': self.DATASET_ID,
+            'namespace': None,
+        }
+        self.assertEqual(key.kwargs, expected_kwargs)
+
+    def test_key_w_namespace(self):
+        from gcloud.datastore import dataset as MUT
+        from gcloud._testing import _Monkey
+
+        KIND = 'KIND'
+        ID = 1234
+        NAMESPACE = object()
+        dataset = self._makeOne(namespace=NAMESPACE)
+        with _Monkey(MUT, Key=_Dummy):
+            key = dataset.key(KIND, ID)
+
+        self.assertTrue(isinstance(key, _Dummy))
+        expected_kwargs = {
+            'dataset_id': self.DATASET_ID,
+            'namespace': NAMESPACE,
+        }
+        self.assertEqual(key.kwargs, expected_kwargs)
+
+    def test_key_w_namespace_collision(self):
+        from gcloud.datastore import dataset as MUT
+        from gcloud._testing import _Monkey
+
+        KIND = 'KIND'
+        ID = 1234
+        NAMESPACE1 = object()
+        NAMESPACE2 = object()
+        dataset = self._makeOne(namespace=NAMESPACE1)
+        with _Monkey(MUT, Key=_Dummy):
+            key = dataset.key(KIND, ID, namespace=NAMESPACE2)
+
+        self.assertTrue(isinstance(key, _Dummy))
+        expected_kwargs = {
+            'dataset_id': self.DATASET_ID,
+            'namespace': NAMESPACE2,
+        }
+        self.assertEqual(key.kwargs, expected_kwargs)
 
     def test_batch_wo_connection(self):
         from gcloud.datastore import dataset as MUT
@@ -347,8 +391,11 @@ def test_transaction_w_connection(self):
 
         self.assertTrue(isinstance(xact, _Dummy))
         self.assertEqual(xact.args, ())
-        self.assertEqual(xact.kwargs,
-                         {'dataset_id': self.DATASET_ID, 'connection': conn})
+        expected_kwargs = {
+            'dataset_id': self.DATASET_ID,
+            'connection': conn,
+        }
+        self.assertEqual(xact.kwargs, expected_kwargs)
 
     def test_query_w_dataset_id(self):
         KIND = 'KIND'
@@ -366,7 +413,11 @@ def test_query_w_defaults(self):
 
         self.assertTrue(isinstance(query, _Dummy))
         self.assertEqual(query.args, ())
-        self.assertEqual(query.kwargs, {'dataset_id': self.DATASET_ID})
+        expected_kwargs = {
+            'dataset_id': self.DATASET_ID,
+            'namespace': None,
+        }
+        self.assertEqual(query.kwargs, expected_kwargs)
 
     def test_query_explicit(self):
         from gcloud.datastore import dataset as MUT
@@ -405,6 +456,45 @@ def test_query_explicit(self):
         self.assertEqual(query.args, ())
         self.assertEqual(query.kwargs, kwargs)
 
+    def test_query_w_namespace(self):
+        from gcloud.datastore import dataset as MUT
+        from gcloud._testing import _Monkey
+
+        KIND = 'KIND'
+        NAMESPACE = object()
+        dataset = self._makeOne(namespace=NAMESPACE)
+        with _Monkey(MUT, Query=_Dummy):
+            query = dataset.query(kind=KIND)
+
+        self.assertTrue(isinstance(query, _Dummy))
+        self.assertEqual(query.args, ())
+        expected_kwargs = {
+            'dataset_id': self.DATASET_ID,
+            'namespace': NAMESPACE,
+            'kind': KIND,
+        }
+        self.assertEqual(query.kwargs, expected_kwargs)
+
+    def test_query_w_namespace_collision(self):
+        from gcloud.datastore import dataset as MUT
+        from gcloud._testing import _Monkey
+
+        KIND = 'KIND'
+        NAMESPACE1 = object()
+        NAMESPACE2 = object()
+        dataset = self._makeOne(namespace=NAMESPACE1)
+        with _Monkey(MUT, Query=_Dummy):
+            query = dataset.query(kind=KIND, namespace=NAMESPACE2)
+
+        self.assertTrue(isinstance(query, _Dummy))
+        self.assertEqual(query.args, ())
+        expected_kwargs = {
+            'dataset_id': self.DATASET_ID,
+            'namespace': NAMESPACE2,
+            'kind': KIND,
+        }
+        self.assertEqual(query.kwargs, expected_kwargs)
+
 
 class _Dummy(object):