add test for entities import/export

datastore/cloud-client/admin.py

@@ -35,10 +35,10 @@ def export_entities(project_id, output_url_prefix):
     client = datastore_admin_client.DatastoreAdminClient()
 
     op = client.export_entities(project_id, output_url_prefix)
-    entities = op.result()
+    response = op.result()
 
     print("Entities were exported\n")
-    return entities
+    return response
 # [END datastore_admin_entities_export]
 
 
@@ -50,9 +50,10 @@ def import_entities(project_id, input_url):
     client = datastore_admin_client.DatastoreAdminClient()
 
     op = client.import_entities(project_id, input_url)
-    op.result()
+    response = op.result()
 
     print("Entities were imported\n")
+    return response
 # [END datastore_admin_entities_import]
 
 

datastore/cloud-client/admin_test.py

@@ -13,10 +13,12 @@
 
 import os
 import pytest
+from retrying import retry
 
 import admin
 
 PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
+BUCKET = os.environ["CLOUD_STORAGE_BUCKET"]
 
 
 class TestDatastoreAdminSnippets:
@@ -35,3 +37,10 @@ def test_get_index(self):
 
     def test_list_index(self):
         assert admin.list_indexes(PROJECT)
+
+    @retry(stop_max_attempt_number=3, stop_max_delay=540000)
+    def test_export_import_entities(self):
+        response = admin.export_entities(PROJECT, "gs://" + BUCKET)
+        assert response
+
+        assert admin.import_entities(PROJECT, response.output_url)