You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2019/12/11 14:35:51 UTC

[GitHub] [airflow] mik-laj commented on a change in pull request #6777: [AIRFLOW-6223] BigQuery - improve hook tests [depends on AIRFLOW-6220]

mik-laj commented on a change in pull request #6777: [AIRFLOW-6223] BigQuery - improve hook tests [depends on AIRFLOW-6220]
URL: https://github.com/apache/airflow/pull/6777#discussion_r356633037
 
 

 ##########
 File path: tests/gcp/hooks/test_bigquery.py
 ##########
 @@ -177,105 +119,146 @@ def test_invalid_syntax(self, table_input, var_name, exception_message):
 
 
 class TestBigQueryHookSourceFormat(unittest.TestCase):
-    def test_invalid_source_format(self):
-        with self.assertRaises(Exception) as context:
-            hook.BigQueryBaseCursor("test", "test").run_load(
+    @mock.patch(
+        'airflow.gcp.hooks.base.CloudBaseHook._get_credentials_and_project_id',
+        return_value=("CREDENTIALS", "PROJECT_ID",)
+    )
+    @mock.patch("airflow.gcp.hooks.bigquery.BigQueryHook.get_service")
+    def test_invalid_source_format(self, mock_get_service, mock_get_creds_and_proj_id):
+        with self.assertRaisesRegex(
+            Exception,
+            r"JSON is not a valid source format. Please use one of the following types: \['CSV', "
+            r"'NEWLINE_DELIMITED_JSON', 'AVRO', 'GOOGLE_SHEETS', 'DATASTORE_BACKUP', 'PARQUET'\]"
+        ):
+            bq_hook = hook.BigQueryHook()
+            cursor = bq_hook.get_cursor()
+            cursor.run_load(
                 "test.test", "test_schema.json", ["test_data.json"], source_format="json"
             )
 
-        # since we passed 'json' in, and it's not valid, make sure it's present in the
-        # error string.
-        self.assertIn("JSON", str(context.exception))
-
 
 class TestBigQueryExternalTableSourceFormat(unittest.TestCase):
-    def test_invalid_source_format(self):
-        with self.assertRaises(Exception) as context:
-            hook.BigQueryBaseCursor("test", "test").create_external_table(
+    @mock.patch(
+        'airflow.gcp.hooks.base.CloudBaseHook._get_credentials_and_project_id',
+        return_value=("CREDENTIALS", "PROJECT_ID",)
+    )
+    @mock.patch("airflow.gcp.hooks.bigquery.BigQueryHook.get_service")
+    def test_invalid_source_format(self, mock_get_service, mock_get_creds_and_proj_id):
+        with self.assertRaisesRegex(
+            Exception,
+            r"JSON is not a valid source format. Please use one of the following types: \['CSV', "
+            r"'NEWLINE_DELIMITED_JSON', 'AVRO', 'GOOGLE_SHEETS', 'DATASTORE_BACKUP', 'PARQUET'\]"
+        ):
+            bq_hook = hook.BigQueryHook()
+            cursor = bq_hook.get_cursor()
+            cursor.create_external_table(
                 external_project_dataset_table='test.test',
                 schema_fields='test_schema.json',
                 source_uris=['test_data.json'],
                 source_format='json'
             )
 
-        # since we passed 'csv' in, and it's not valid, make sure it's present in the
-        # error string.
-        self.assertIn("JSON", str(context.exception))
-
-
-# Helpers to test_cancel_queries that have mock_poll_job_complete returning false,
-# unless mock_job_cancel was called with the same job_id
-mock_canceled_jobs = []  # type: List
-
-
-def mock_poll_job_complete(job_id):
-    return job_id in mock_canceled_jobs
-
-
-# pylint: disable=invalid-name
-# noinspection PyUnusedLocal
-def mock_job_cancel(projectId, jobId):  # pylint: disable=unused-argument
-    mock_canceled_jobs.append(jobId)
-    return mock.Mock()
-
 
 class TestBigQueryBaseCursor(unittest.TestCase):
-    def test_invalid_schema_update_options(self):
-        with self.assertRaises(Exception) as context:
-            hook.BigQueryBaseCursor("test", "test").run_load(
+    @mock.patch(
+        'airflow.gcp.hooks.base.CloudBaseHook._get_credentials_and_project_id',
+        return_value=("CREDENTIALS", "PROJECT_ID",)
+    )
+    @mock.patch("airflow.gcp.hooks.bigquery.BigQueryHook.get_service")
+    def test_invalid_schema_update_options(self, mock_get_service, mock_get_creds_and_proj_id):
+        with self.assertRaisesRegex(
+            Exception,
+            r"\['THIS IS NOT VALID'\] contains invalid schema update options.Please only use one or more of "
+            r"the following options: \['ALLOW_FIELD_ADDITION', 'ALLOW_FIELD_RELAXATION'\]"
+        ):
+            bq_hook = hook.BigQueryHook()
+            cursor = bq_hook.get_cursor()
+            cursor.run_load(
                 "test.test",
                 "test_schema.json",
                 ["test_data.json"],
                 schema_update_options=["THIS IS NOT VALID"]
             )
-        self.assertIn("THIS IS NOT VALID", str(context.exception))
 
-    def test_invalid_schema_update_and_write_disposition(self):
-        with self.assertRaises(Exception) as context:
-            hook.BigQueryBaseCursor("test", "test").run_load(
+    @mock.patch(
+        'airflow.gcp.hooks.base.CloudBaseHook._get_credentials_and_project_id',
+        return_value=("CREDENTIALS", "PROJECT_ID",)
+    )
+    @mock.patch("airflow.gcp.hooks.bigquery.BigQueryHook.get_service")
+    def test_invalid_schema_update_and_write_disposition(self, mock_get_service, mock_get_creds_and_proj_id):
+        with self.assertRaisesRegex(Exception, "schema_update_options is only allowed if"
+                                               " write_disposition is 'WRITE_APPEND' or 'WRITE_TRUNCATE'."):
+            bq_hook = hook.BigQueryHook()
+            cursor = bq_hook.get_cursor()
+            cursor.run_load(
                 "test.test",
                 "test_schema.json",
                 ["test_data.json"],
                 schema_update_options=['ALLOW_FIELD_ADDITION'],
                 write_disposition='WRITE_EMPTY'
             )
-        self.assertIn("schema_update_options is only", str(context.exception))
 
-    def test_cancel_queries(self):
+    @mock.patch("airflow.gcp.hooks.bigquery.BigQueryBaseCursor.poll_job_complete", side_effect=[False, True])
+    @mock.patch('airflow.gcp.hooks.base.CloudBaseHook._get_credentials_and_project_id')
 
 Review comment:
   ```suggestion
       @mock.patch('airflow.gcp.hooks.base.CloudBaseHook._get_credentials_and_project_id', return_value=("CREDENTIALS", "PROJECT_ID",))
   ```

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services