fix(datastore): properly retry a flaky test (GoogleCloudPlatform#5543)

fixes GoogleCloudPlatform#5495

Author: Takashi Matsuo

datastore/cloud-client/admin_test.py

@@ -11,6 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from concurrent.futures import TimeoutError
 import os
 
 import backoff
@@ -41,7 +42,7 @@ def test_get_index(self):
     def test_list_index(self):
         assert admin.list_indexes(PROJECT)
 
-    @backoff.on_exception(backoff.expo, RetryError, max_tries=3)
+    @backoff.on_exception(backoff.expo, (RetryError, TimeoutError), max_tries=3)
     def test_export_import_entities(self):
         response = admin.export_entities(PROJECT, "gs://" + BUCKET)
         assert response

datastore/cloud-client/snippets_test.py

@@ -31,7 +31,7 @@ def __init__(self, *args, **kwargs):
     def cleanup(self):
         with self.batch():
             self.delete_multi(
-                list(set([x.key for x in self.entities_to_delete]))
+                list(set([x.key for x in self.entities_to_delete if x]))
                 + list(set(self.keys_to_delete))
             )