[AIRFLOW-6970] Improve GCP Video Intelligence system tests (#7604)

apache · Mar 3, 2020 · 0d1e308 · 0d1e308
1 parent 46e55fb
commit 0d1e308
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 12 deletions.
diff --git a/airflow/providers/google/cloud/hooks/base.py b/airflow/providers/google/cloud/hooks/base.py
@@ -34,7 +34,7 @@
 import httplib2
 import tenacity
 from google.api_core.exceptions import (
-    AlreadyExists, Forbidden, GoogleAPICallError, ResourceExhausted, RetryError,
+    AlreadyExists, Forbidden, GoogleAPICallError, ResourceExhausted, RetryError, TooManyRequests,
 )
 from google.api_core.gapic_v1.client_info import ClientInfo
 from google.auth.environment_vars import CREDENTIALS
@@ -76,12 +76,12 @@ def is_soft_quota_exception(exception: Exception):
     """
     if isinstance(exception, Forbidden):
         return any(
-            reason in error["reason"]
+            reason in error.details()
             for reason in INVALID_REASONS
             for error in exception.errors
         )
 
-    if isinstance(exception, ResourceExhausted):
+    if isinstance(exception, (ResourceExhausted, TooManyRequests)):
         return any(
             key in error.details()
             for key in INVALID_KEYS

diff --git a/tests/providers/google/cloud/hooks/test_base.py b/tests/providers/google/cloud/hooks/test_base.py
@@ -81,11 +81,7 @@ def test_do_nothing_on_non_error(self):
     def test_retry_on_exception(self):
         message = "POST https://translation.googleapis.com/language/translate/v2: User Rate Limit Exceeded"
         errors = [
-            {
-                'message': 'User Rate Limit Exceeded',
-                'domain': 'usageLimits',
-                'reason': 'userRateLimitExceeded',
-            }
+            mock.MagicMock(details=mock.PropertyMock(return_value='userRateLimitExceeded'))
         ]
         custom_fn = NoForbiddenAfterCount(
             count=5,
@@ -99,7 +95,7 @@ def test_raise_exception_on_non_quota_exception(self):
         with self.assertRaisesRegex(Forbidden, "Daily Limit Exceeded"):
             message = "POST https://translation.googleapis.com/language/translate/v2: Daily Limit Exceeded"
             errors = [
-                {'message': 'Daily Limit Exceeded', 'domain': 'usageLimits', 'reason': 'dailyLimitExceeded'}
+                mock.MagicMock(details=mock.PropertyMock(return_value='dailyLimitExceeded'))
             ]
 
             _retryable_test_with_temporary_quota_retry(

diff --git a/tests/providers/google/cloud/operators/test_video_intelligence_system.py b/tests/providers/google/cloud/operators/test_video_intelligence_system.py
@@ -19,12 +19,12 @@
 
 import pytest
 
+from airflow.providers.google.cloud.example_dags.example_video_intelligence import GCP_BUCKET_NAME
 from tests.providers.google.cloud.utils.gcp_authenticator import GCP_AI_KEY, GCP_GCS_KEY
 from tests.test_utils.gcp_system_helpers import CLOUD_DAG_FOLDER, GoogleSystemTest, provide_gcp_context
 
 GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "example-project")
-GCP_BUCKET_NAME = os.environ.get("GCP_VIDEO_INTELLIGENCE_BUCKET_NAME", "test-bucket-name")
-GCP_VIDEO_SOURCE_URL = os.environ.get("GCP_VIDEO_INTELLIGENCE_VIDEO_SOURCE_URL", "http://nasa.gov")
+GCP_VIDEO_SOURCE_URL = "https://www.sample-videos.com/video123/mp4/720/big_buck_bunny_720p_1mb.mp4"
 
 
 @pytest.mark.backend("mysql", "postgres")
@@ -49,5 +49,5 @@ def tearDown(self):
         super().tearDown()
 
     @provide_gcp_context(GCP_AI_KEY)
-    def test_run_example_dag_spanner(self):
+    def test_example_dag(self):
         self.run_dag('example_gcp_video_intelligence', CLOUD_DAG_FOLDER)