def test_credentials_provided(self, mock_read_gbq, mock_get_credentials_and_project_id):
self.instance = hook.BigQueryHook()
self.instance.get_pandas_df('select 1')
args, kwargs = mock_read_gbq.call_args
self.assertEqual("CREDENTIALS", kwargs['credentials'])
self.assertEqual("PROJECT_ID", kwargs['project_id'])
def test_location_propagates_properly(self, run_with_config):
with mock.patch.object(hook.BigQueryHook, 'get_service'):
bq_hook = hook.BigQueryHook(location=None)
self.assertIsNone(bq_hook.location)
bq_cursor = hook.BigQueryBaseCursor(mock.Mock(),
'test-project',
location=None)
self.assertIsNone(bq_cursor.location)
bq_cursor.run_query(sql='select 1', location='US')
assert run_with_config.call_count == 1
self.assertEqual(bq_cursor.location, 'US')
def setUp(self):
self.instance = hook.BigQueryHook()
def setUp(self):
self.hook = hook.BigQueryHook()
def test_legacy_sql_override_propagates_properly(self, run_with_config):
with mock.patch.object(hook.BigQueryHook, 'get_service'):
bq_hook = hook.BigQueryHook(use_legacy_sql=False)
bq_hook.get_first('query')
args, kwargs = run_with_config.call_args
self.assertIs(args[0]['query']['useLegacySql'], False)
def test_hook_uses_legacy_sql_by_default(self, run_with_config):
with mock.patch.object(hook.BigQueryHook, 'get_service'):
bq_hook = hook.BigQueryHook()
bq_hook.get_first('query')
args, kwargs = run_with_config.call_args
self.assertIs(args[0]['query']['useLegacySql'], True)
import unittest
from unittest import mock
from typing import List, Optional
from google.auth.exceptions import GoogleAuthError
from googleapiclient.errors import HttpError
from airflow.gcp.hooks import bigquery as hook
from airflow.gcp.hooks.bigquery import _cleanse_time_partitioning, \
_validate_value, _api_resource_configs_duplication_check, \
_validate_src_fmt_configs
bq_available = True
try:
hook.BigQueryHook().get_service()
except GoogleAuthError:
bq_available = False
class TestBigQueryHookConnection(unittest.TestCase):
hook = None # type: Optional[hook.BigQueryHook]
def setUp(self):
self.hook = hook.BigQueryHook()
@mock.patch("airflow.gcp.hooks.bigquery.BigQueryConnection")
@mock.patch("airflow.gcp.hooks.bigquery.BigQueryHook._authorize")
@mock.patch("airflow.gcp.hooks.bigquery.build")
def test_bigquery_client_creation(self, mock_build, mock_authorize, mock_bigquery_connection):
result = self.hook.get_conn()