From 2aea38361e28a5cac213661ed29dd3bc9a9a7f4d Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Tue, 22 Nov 2022 17:35:15 -0500 Subject: [PATCH] Additional unit tests for BigQuery Connector (#746) I wrote unit tests to try to programmatically reproduce the issue described in #685. However, I was unable to identify the issue. --- parsons/google/utitities.py | 7 ++- test/test_google/test_google_bigquery.py | 51 +++++++++++++++- test/test_google/test_utilities.py | 74 ++++++++++++++++++++++++ 3 files changed, 127 insertions(+), 5 deletions(-) create mode 100644 test/test_google/test_utilities.py diff --git a/parsons/google/utitities.py b/parsons/google/utitities.py index 7f29a09554..196a3d13c1 100644 --- a/parsons/google/utitities.py +++ b/parsons/google/utitities.py @@ -1,16 +1,19 @@ +import typing as t from parsons.utilities import files from parsons.utilities import check_env import json import os -def setup_google_application_credentials(app_creds, env_var_name='GOOGLE_APPLICATION_CREDENTIALS'): +def setup_google_application_credentials( + app_creds: t.Union[t.Dict, str, None], + env_var_name: str = 'GOOGLE_APPLICATION_CREDENTIALS') -> None: # Detect if app_creds is a dict, path string or json string, and if it is a # json string, then convert it to a temporary file. Then set the # environmental variable. credentials = check_env.check(env_var_name, app_creds) try: - if (type(credentials) is dict): + if type(credentials) is dict: credentials = json.dumps(credentials) if json.loads(credentials): creds_path = files.string_to_temp_file(credentials, suffix='.json') diff --git a/test/test_google/test_google_bigquery.py b/test/test_google/test_google_bigquery.py index f7a7e9f88a..e6d47c470f 100644 --- a/test/test_google/test_google_bigquery.py +++ b/test/test_google/test_google_bigquery.py @@ -1,16 +1,45 @@ +import json import os -import unittest import unittest.mock as mock + from google.cloud import bigquery from google.cloud import exceptions + from parsons import GoogleBigQuery, Table +from parsons.google.google_cloud_storage import GoogleCloudStorage +from test.test_google.test_utilities import FakeCredentialTest + + +class FakeClient: + """A Fake Storage Client used for monkey-patching.""" + def __init__(self, project=None): + self.project = project + + +class FakeGoogleCloudStorage(GoogleCloudStorage): + """A Fake GoogleCloudStorage object used to test setting up credentials.""" + + @mock.patch('google.cloud.storage.Client', FakeClient) + def __init__(self): + super().__init__(None, None) + + def upload_table(self, table, bucket_name, blob_name, data_type='csv', default_acl=None): + pass + + def delete_blob(self, bucket_name, blob_name): + pass -class TestGoogleBigQuery(unittest.TestCase): +class TestGoogleBigQuery(FakeCredentialTest): def setUp(self): - os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = 'foo' + super().setUp() + os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = self.cred_path self.tmp_gcs_bucket = 'tmp' + def tearDown(self) -> None: + super().tearDown() + del os.environ['GOOGLE_APPLICATION_CREDENTIALS'] + def test_query(self): query_string = 'select * from table' @@ -142,6 +171,22 @@ def test_copy__bad_if_exists(self): bq.copy(self.default_table, 'dataset.table', tmp_gcs_bucket=self.tmp_gcs_bucket, if_exists='foo', gcs_client=gcs_client) + def test_copy__credentials_are_correctly_set(self): + tbl = self.default_table + bq = self._build_mock_client_for_copying(table_exists=False) + + # Pass in our fake GCS Client. + bq.copy(tbl, 'dataset.table', tmp_gcs_bucket=self.tmp_gcs_bucket, + gcs_client=FakeGoogleCloudStorage()) + + actual = os.environ['GOOGLE_APPLICATION_CREDENTIALS'] + + with open(actual, 'r') as factual: + with open(self.cred_path, 'r') as fexpected: + actual_str = factual.read() + self.assertEqual(actual_str, fexpected.read()) + self.assertEqual(self.cred_contents, json.loads(actual_str)) + def _build_mock_client_for_querying(self, results): # Create a mock that will play the role of the cursor cursor = mock.MagicMock() diff --git a/test/test_google/test_utilities.py b/test/test_google/test_utilities.py new file mode 100644 index 0000000000..0b33d013be --- /dev/null +++ b/test/test_google/test_utilities.py @@ -0,0 +1,74 @@ +import json +import unittest +import os +import tempfile + +from parsons.google import utitities as util + + +class FakeCredentialTest(unittest.TestCase): + def setUp(self) -> None: + self.dir = tempfile.TemporaryDirectory() + self.cred_path = os.path.join(self.dir.name, 'mycred.json') + self.cred_contents = { + "client_id": "foobar.apps.googleusercontent.com", + "client_secret": str(hash("foobar")), + "quota_project_id": "project-id", + "refresh_token": str(hash("foobarfoobar")), + "type": "authorized_user" + } + with open(self.cred_path, 'w') as f: + json.dump(self.cred_contents, f) + + def tearDown(self) -> None: + self.dir.cleanup() + + +class TestSetupGoogleApplicationCredentials(FakeCredentialTest): + TEST_ENV_NAME = 'DUMMY_APP_CREDS' + + def tearDown(self) -> None: + super().tearDown() + del os.environ[self.TEST_ENV_NAME] + + def test_noop_if_env_already_set(self): + os.environ[self.TEST_ENV_NAME] = self.cred_path + util.setup_google_application_credentials(None, self.TEST_ENV_NAME) + self.assertEqual(os.environ[self.TEST_ENV_NAME], self.cred_path) + + def test_accepts_dictionary(self): + util.setup_google_application_credentials(self.cred_contents, self.TEST_ENV_NAME) + actual = os.environ[self.TEST_ENV_NAME] + self.assertTrue(os.path.exists(actual)) + with open(actual, 'r') as f: + self.assertEqual(json.load(f), self.cred_contents) + + def test_accepts_string(self): + cred_str = json.dumps(self.cred_contents) + util.setup_google_application_credentials(cred_str, self.TEST_ENV_NAME) + actual = os.environ[self.TEST_ENV_NAME] + self.assertTrue(os.path.exists(actual)) + with open(actual, 'r') as f: + self.assertEqual(json.load(f), self.cred_contents) + + def test_accepts_file_path(self): + util.setup_google_application_credentials(self.cred_path, self.TEST_ENV_NAME) + actual = os.environ[self.TEST_ENV_NAME] + self.assertTrue(os.path.exists(actual)) + with open(actual, 'r') as f: + self.assertEqual(json.load(f), self.cred_contents) + + def test_credentials_are_valid_after_double_call(self): + # write creds to tmp file... + util.setup_google_application_credentials(self.cred_contents, self.TEST_ENV_NAME) + fst = os.environ[self.TEST_ENV_NAME] + + # repeat w/ default args... + util.setup_google_application_credentials(None, self.TEST_ENV_NAME) + snd = os.environ[self.TEST_ENV_NAME] + + with open(fst, 'r') as ffst: + with open(snd, 'r') as fsnd: + actual = fsnd.read() + self.assertEqual(self.cred_contents, json.loads(actual)) + self.assertEqual(ffst.read(), actual)