From 3f3bd1b73f81fdcbc78cf87ee3e7f82649fe7095 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Fri, 26 Jan 2024 18:54:30 +0000 Subject: [PATCH 1/3] test: fix unit tests for handling drive access issues --- tests/unit/session/test_session.py | 37 ++++++++++++++++++++++++------ 1 file changed, 30 insertions(+), 7 deletions(-) diff --git a/tests/unit/session/test_session.py b/tests/unit/session/test_session.py index 5fc8996993..33357af90a 100644 --- a/tests/unit/session/test_session.py +++ b/tests/unit/session/test_session.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy +import datetime import os from unittest import mock @@ -63,15 +65,36 @@ def test_read_gbq_not_found_tables(not_found_table_id): ], ) def test_read_gbq_external_table_no_drive_access(api_name, query_or_table): - bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True) - bqclient.project = "test-project" - bqclient.get_table.side_effect = google.api_core.exceptions.Forbidden( - "Access Denied: BigQuery BigQuery: Permission denied while getting Drive credentials." - ) - session = resources.create_bigquery_session(bqclient=bqclient) + session = resources.create_bigquery_session() + mock_query_job = session.bqclient.query.return_value + + def query_mock(query, *args, **kwargs): + query_job = copy.copy(mock_query_job) + + if query.lstrip().startswith("SELECT *"): + raise google.api_core.exceptions.Forbidden( + "Access Denied: BigQuery BigQuery: Permission denied while getting Drive credentials." + ) + + if query.startswith("SELECT CURRENT_TIMESTAMP()"): + query_job.result = mock.MagicMock(return_value=[[datetime.datetime.now()]]) + + return query_job + + session.bqclient.query = query_mock + + def get_table_mock(dataset_ref): + dataset = google.cloud.bigquery.Dataset(dataset_ref) + dataset.location = session._location + return dataset + + session.bqclient.get_table = get_table_mock api = getattr(session, api_name) - with pytest.raises(google.api_core.exceptions.Forbidden): + with pytest.raises( + google.api_core.exceptions.Forbidden, + match="Check https://cloud.google.com/bigquery/docs/query-drive-data#Google_Drive_permissions.", + ): api(query_or_table) From f7f2bbbd40ba6dcb3dbdc593eb6c08512556c34a Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Mon, 5 Feb 2024 22:29:53 +0000 Subject: [PATCH 2/3] refactor test logic --- tests/unit/resources.py | 24 ++++++++++++++++-------- tests/unit/session/test_session.py | 11 ++--------- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/tests/unit/resources.py b/tests/unit/resources.py index b239b04671..967e42548f 100644 --- a/tests/unit/resources.py +++ b/tests/unit/resources.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import datetime from typing import Dict, List, Optional import unittest.mock as mock @@ -49,14 +50,21 @@ def create_bigquery_session( "test_dataset", ) - query_job = mock.create_autospec(google.cloud.bigquery.QueryJob) - type(query_job).destination = mock.PropertyMock( - return_value=anonymous_dataset.table("test_table"), - ) - type(query_job).session_info = google.cloud.bigquery.SessionInfo( - {"sessionInfo": {"sessionId": session_id}}, - ) - bqclient.query.return_value = query_job + def query_mock(query, *args, **kwargs): + query_job = mock.create_autospec(google.cloud.bigquery.QueryJob) + type(query_job).destination = mock.PropertyMock( + return_value=anonymous_dataset.table("test_table"), + ) + type(query_job).session_info = google.cloud.bigquery.SessionInfo( + {"sessionInfo": {"sessionId": session_id}}, + ) + + if query.startswith("SELECT CURRENT_TIMESTAMP()"): + query_job.result = mock.MagicMock(return_value=[[datetime.datetime.now()]]) + + return query_job + + bqclient.query = query_mock clients_provider = mock.create_autospec(bigframes.session.clients.ClientsProvider) type(clients_provider).bqclient = mock.PropertyMock(return_value=bqclient) diff --git a/tests/unit/session/test_session.py b/tests/unit/session/test_session.py index 33357af90a..ea8d0882ae 100644 --- a/tests/unit/session/test_session.py +++ b/tests/unit/session/test_session.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy -import datetime import os from unittest import mock @@ -66,20 +64,15 @@ def test_read_gbq_not_found_tables(not_found_table_id): ) def test_read_gbq_external_table_no_drive_access(api_name, query_or_table): session = resources.create_bigquery_session() - mock_query_job = session.bqclient.query.return_value + session_query_mock = session.bqclient.query def query_mock(query, *args, **kwargs): - query_job = copy.copy(mock_query_job) - if query.lstrip().startswith("SELECT *"): raise google.api_core.exceptions.Forbidden( "Access Denied: BigQuery BigQuery: Permission denied while getting Drive credentials." ) - if query.startswith("SELECT CURRENT_TIMESTAMP()"): - query_job.result = mock.MagicMock(return_value=[[datetime.datetime.now()]]) - - return query_job + return session_query_mock(query, *args, **kwargs) session.bqclient.query = query_mock From 8d8496eabc873c7a28ee98bd1c1901a90ee21edb Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Tue, 6 Feb 2024 01:20:44 +0000 Subject: [PATCH 3/3] adapt new create_bigquery_session in one failing test --- tests/unit/test_compute_options.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_compute_options.py b/tests/unit/test_compute_options.py index a613bca7b9..2de715a40e 100644 --- a/tests/unit/test_compute_options.py +++ b/tests/unit/test_compute_options.py @@ -11,6 +11,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from unittest import mock + import bigframes as bf from . import resources @@ -18,8 +20,8 @@ def test_maximum_bytes_option(): session = resources.create_bigquery_session() + session.bqclient.query = mock.MagicMock() with bf.option_context("compute.maximum_bytes_billed", 10000): - session.bqclient.query.reset_mock() session._start_query("query") call = session.bqclient.query.call_args assert call.kwargs["job_config"].maximum_bytes_billed == 10000