fix(datastore): properly retry the test (GoogleCloudPlatform#5496)

Takashi Matsuo · web-flow · commit 38b406e1deaf · 2021-03-09T09:48:21.000-08:00
fixes GoogleCloudPlatform#5495
diff --git a/datastore/cloud-client/admin_test.py b/datastore/cloud-client/admin_test.py
@@ -14,6 +14,7 @@
 import os
 
 import backoff
+from google.api_core.exceptions import RetryError
 import pytest
 
 import admin
@@ -40,7 +41,7 @@ def test_get_index(self):
     def test_list_index(self):
         assert admin.list_indexes(PROJECT)
 
-    @backoff.on_exception(backoff.expo, AssertionError, max_tries=3)
+    @backoff.on_exception(backoff.expo, RetryError, max_tries=3)
     def test_export_import_entities(self):
         response = admin.export_entities(PROJECT, "gs://" + BUCKET)
         assert response
