Merge pull request #980 from dhermes/fix-datastore-plumbing-modules

Fix datastore test plumbing modules.
googleapis · Jul 11, 2015 · c8593bf · c8593bf
2 parents 3c76dbf + c190be4
commit c8593bf
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/system_tests/clear_datastore.py b/system_tests/clear_datastore.py
@@ -21,6 +21,7 @@
 
 
 _implicit_environ._DATASET_ENV_VAR_NAME = 'GCLOUD_TESTS_DATASET_ID'
+CLIENT = datastore.Client()
 
 
 FETCH_MAX = 20
@@ -36,7 +37,7 @@
 
 def fetch_keys(kind, fetch_max=FETCH_MAX, query=None, cursor=None):
     if query is None:
-        query = datastore.Query(kind=kind, projection=['__key__'])
+        query = CLIENT.query(kind=kind, projection=['__key__'])
 
     iterator = query.fetch(limit=fetch_max, start_cursor=cursor)
 
@@ -65,7 +66,7 @@ def remove_kind(kind):
         return
 
     delete_outside_transaction = False
-    with datastore.Transaction():
+    with CLIENT.transaction():
         # Now that we have all results, we seek to delete.
         print('Deleting keys:')
         print(results)
@@ -74,10 +75,10 @@ def remove_kind(kind):
         if len(ancestors) > TRANSACTION_MAX_GROUPS:
             delete_outside_transaction = True
         else:
-            datastore.delete_multi([result.key for result in results])
+            CLIENT.delete_multi([result.key for result in results])
 
     if delete_outside_transaction:
-        datastore.delete_multi([result.key for result in results])
+        CLIENT.delete_multi([result.key for result in results])
 
 
 def remove_all_entities():

diff --git a/system_tests/populate_datastore.py b/system_tests/populate_datastore.py
@@ -21,6 +21,7 @@
 
 
 _implicit_environ._DATASET_ENV_VAR_NAME = 'GCLOUD_TESTS_DATASET_ID'
+CLIENT = datastore.Client()
 
 
 ANCESTOR = ('Book', 'GoT')
@@ -82,12 +83,12 @@
 
 
 def add_characters():
-    with datastore.Transaction() as xact:
+    with CLIENT.transaction() as xact:
         for key_path, character in zip(KEY_PATHS, CHARACTERS):
             if key_path[-1] != character['name']:
                 raise ValueError(('Character and key don\'t agree',
                                   key_path, character))
-            entity = datastore.Entity(key=datastore.Key(*key_path))
+            entity = datastore.Entity(key=CLIENT.key(*key_path))
             entity.update(character)
             xact.put(entity)
             print('Adding Character %s %s' % (character['name'],