test: fix unit tests for handling drive access issues

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)

tests/unit/session/test_session.py

@@ -63,15 +63,31 @@ 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()
+    session_query_mock = session.bqclient.query
+
+    def query_mock(query, *args, **kwargs):
+        if query.lstrip().startswith("SELECT *"):
+            raise google.api_core.exceptions.Forbidden(
+                "Access Denied: BigQuery BigQuery: Permission denied while getting Drive credentials."
+            )
+
+        return session_query_mock(query, *args, **kwargs)
+
+    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)
 
 

tests/unit/test_compute_options.py

@@ -11,15 +11,17 @@
 # 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
 
 
 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