test: fix unit test coverage for google-cloud-bigquery

Author: shuoweil

packages/google-cloud-bigquery/.coveragerc

@@ -12,3 +12,5 @@ exclude_lines =
     pragma: (no cover|NO COVER)
     # Ignore debug-only repr
     def __repr__
+    except ImportError:
+    except ImportError as .*:

packages/google-cloud-bigquery/google/cloud/bigquery/magics/magics.py

@@ -23,10 +23,10 @@
 
 from __future__ import print_function
 
-import re
 import ast
 import copy
 import functools
+import re
 import sys
 import time
 import warnings
@@ -39,14 +39,12 @@
 except ImportError:
     raise ImportError("This module can only be loaded in IPython.")
 
-from google.api_core import client_info
-from google.api_core import client_options
-from google.api_core.exceptions import NotFound
 import google.auth  # type: ignore
-from google.cloud import bigquery
 import google.cloud.bigquery.dataset
-from google.cloud.bigquery import _versions_helpers
-from google.cloud.bigquery import exceptions
+from google.api_core import client_info, client_options
+from google.api_core.exceptions import NotFound
+from google.cloud import bigquery
+from google.cloud.bigquery import _versions_helpers, exceptions
 from google.cloud.bigquery.dbapi import _helpers
 from google.cloud.bigquery.magics import line_arg_parser as lap
 
@@ -231,7 +229,7 @@ def progress_bar_type(self, value):
 # their code.
 if bigquery_magics is not None:
     context = bigquery_magics.context
-else:
+else:  # pragma: NO COVER
     context = Context()
 
 

packages/google-cloud-bigquery/tests/unit/test__pandas_helpers.py

@@ -18,14 +18,13 @@
 import decimal
 import functools
 import gc
+import importlib.metadata as metadata
 import operator
 import queue
 import time
+import warnings
 from typing import Union
 from unittest import mock
-import warnings
-
-import importlib.metadata as metadata
 
 try:
     import pandas
@@ -45,13 +44,13 @@
     geopandas = None
 
 import pytest
-
 from google import api_core
-
-from google.cloud.bigquery import exceptions
-from google.cloud.bigquery import _pyarrow_helpers
-from google.cloud.bigquery import _versions_helpers
-from google.cloud.bigquery import schema
+from google.cloud.bigquery import (
+    _pyarrow_helpers,
+    _versions_helpers,
+    exceptions,
+    schema,
+)
 from google.cloud.bigquery._pandas_helpers import determine_requested_streams
 
 pyarrow = _versions_helpers.PYARROW_VERSIONS.try_import()
@@ -1831,8 +1830,7 @@ def test__download_table_bqstorage(
     expected_call_count,
     expected_maxsize,
 ):
-    from google.cloud.bigquery import dataset
-    from google.cloud.bigquery import table
+    from google.cloud.bigquery import dataset, table
 
     queue_used = None  # A reference to the queue used by code under test.
 
@@ -1885,10 +1883,9 @@ def test__download_table_bqstorage_shuts_down_workers(
     the child threads are also stopped.
     """
     pytest.importorskip("google.cloud.bigquery_storage_v1")
-    from google.cloud.bigquery import dataset
-    from google.cloud.bigquery import table
     import google.cloud.bigquery_storage_v1.reader
     import google.cloud.bigquery_storage_v1.types
+    from google.cloud.bigquery import dataset, table
 
     monkeypatch.setattr(
         _versions_helpers.BQ_STORAGE_VERSIONS, "_installed_version", None
@@ -2211,10 +2208,10 @@ def test_determine_requested_streams_invalid_max_stream_count():
     bigquery_storage is None, reason="Requires google-cloud-bigquery-storage"
 )
 def test__download_table_bqstorage_w_timeout_error(module_under_test):
-    from google.cloud.bigquery import dataset
-    from google.cloud.bigquery import table
     from unittest import mock
 
+    from google.cloud.bigquery import dataset, table
+
     mock_bqstorage_client = mock.create_autospec(
         bigquery_storage.BigQueryReadClient, instance=True
     )
@@ -2248,10 +2245,10 @@ def slow_download_stream(
     bigquery_storage is None, reason="Requires google-cloud-bigquery-storage"
 )
 def test__download_table_bqstorage_w_timeout_success(module_under_test):
-    from google.cloud.bigquery import dataset
-    from google.cloud.bigquery import table
     from unittest import mock
 
+    from google.cloud.bigquery import dataset, table
+
     mock_bqstorage_client = mock.create_autospec(
         bigquery_storage.BigQueryReadClient, instance=True
     )
@@ -2409,3 +2406,38 @@ def test_download_arrow_bqstorage_passes_timeout_to_create_read_session(
     assert retry_policy is not None
     # Check if deadline is set correctly in the retry policy
     assert retry_policy._deadline == timeout
+
+
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+def test_dataframe_to_bq_schema_w_unused_schema_field(module_under_test):
+    with mock.patch.object(module_under_test, "pandas_gbq", None):
+        with pytest.raises(
+            ValueError, match="bq_schema contains fields not present in dataframe"
+        ):
+            module_under_test.dataframe_to_bq_schema(
+                pandas.DataFrame(), (schema.SchemaField("not_in_df", "STRING"),)
+            )
+
+
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
+def test_get_schema_by_pyarrow_bignumeric(module_under_test):
+    series = pandas.Series([decimal.Decimal("1.12345678901")])
+    result = module_under_test._get_schema_by_pyarrow("col", series)
+    assert result is not None
+    assert result.field_type == "BIGNUMERIC"
+
+
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
+def test_get_types_mapper_range_timestamp_mismatch(module_under_test):
+    if not hasattr(pandas, "ArrowDtype"):
+        return
+    range_ts = pandas.ArrowDtype(
+        pyarrow.struct(
+            [("start", pyarrow.timestamp("us")), ("end", pyarrow.timestamp("us"))]
+        )
+    )
+    mapper = module_under_test.default_types_mapper(range_timestamp_dtype=range_ts)
+    unmatched_struct = pyarrow.struct([("other", pyarrow.int64())])
+    assert mapper(unmatched_struct) is None

packages/google-cloud-bigquery/tests/unit/test__pyarrow_helpers.py

@@ -44,3 +44,16 @@ def test_bq_to_arrow_scalars(module_under_test):
 def test_arrow_scalar_ids_to_bq(module_under_test):
     assert module_under_test.arrow_scalar_ids_to_bq(pyarrow.bool_().id) == "BOOL"
     assert module_under_test.arrow_scalar_ids_to_bq("UNKNOWN_TYPE") is None
+
+
+def test_pyarrow_helpers_when_pyarrow_none(module_under_test):
+    import importlib
+    import sys
+    from unittest import mock
+
+    with mock.patch.dict(sys.modules, {"pyarrow": None}):
+        importlib.reload(module_under_test)
+        assert module_under_test.pyarrow is None
+        assert module_under_test.arrow_scalar_ids_to_bq(1) is None
+
+    importlib.reload(module_under_test)

packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py

@@ -58,14 +58,14 @@ def test_try_import_raises_error_w_legacy_pyarrow():
             versions.try_import(raise_if_error=True)
 
 
-@pytest.mark.skipif(
-    pyarrow is not None,
-    reason="pyarrow is installed, but this test needs it not to be",
-)
 def test_try_import_raises_error_w_no_pyarrow():
+    import sys
+
     versions = _versions_helpers.PyarrowVersions()
-    with pytest.raises(exceptions.LegacyPyarrowError):
-        versions.try_import(raise_if_error=True)
+    with mock.patch.dict(sys.modules, {"pyarrow": None}):
+        assert versions.try_import(raise_if_error=False) is None
+        with pytest.raises(exceptions.LegacyPyarrowError):
+            versions.try_import(raise_if_error=True)
 
 
 @pytest.mark.skipif(pyarrow is None, reason="pyarrow is not installed")
@@ -121,17 +121,29 @@ def test_returns_none_with_legacy_bqstorage():
         assert bq_storage is None
 
 
-@pytest.mark.skipif(
-    bigquery_storage is not None,
-    reason="Tests behavior when `google-cloud-bigquery-storage` isn't installed",
-)
 def test_returns_none_with_bqstorage_uninstalled():
-    try:
-        bqstorage_versions = _versions_helpers.BQStorageVersions()
-        bq_storage = bqstorage_versions.try_import()
-    except exceptions.LegacyBigQueryStorageError:  # pragma: NO COVER
-        raise ("NotFound error raised when raise_if_error == False.")
-    assert bq_storage is None
+    import sys
+
+    from google import cloud
+
+    versions = _versions_helpers.BQStorageVersions()
+    with mock.patch.dict(sys.modules, {"google.cloud.bigquery_storage": None}):
+        with mock.patch.dict(cloud.__dict__):
+            cloud.__dict__.pop("bigquery_storage", None)
+            assert versions.try_import() is None
+
+
+def test_raises_error_with_bqstorage_uninstalled():
+    import sys
+
+    from google import cloud
+
+    versions = _versions_helpers.BQStorageVersions()
+    with mock.patch.dict(sys.modules, {"google.cloud.bigquery_storage": None}):
+        with mock.patch.dict(cloud.__dict__):
+            cloud.__dict__.pop("bigquery_storage", None)
+            with pytest.raises(exceptions.BigQueryStorageNotFoundError):
+                versions.try_import(raise_if_error=True)
 
 
 @pytest.mark.skipif(
@@ -219,14 +231,14 @@ def test_try_import_raises_error_w_legacy_pandas():
             versions.try_import(raise_if_error=True)
 
 
-@pytest.mark.skipif(
-    pandas is not None,
-    reason="pandas is installed, but this test needs it not to be",
-)
 def test_try_import_raises_error_w_no_pandas():
+    import sys
+
     versions = _versions_helpers.PandasVersions()
-    with pytest.raises(exceptions.LegacyPandasError):
-        versions.try_import(raise_if_error=True)
+    with mock.patch.dict(sys.modules, {"pandas": None}):
+        assert versions.try_import(raise_if_error=False) is None
+        with pytest.raises(exceptions.LegacyPandasError):
+            versions.try_import(raise_if_error=True)
 
 
 @pytest.mark.skipif(pandas is None, reason="pandas is not installed")
@@ -255,6 +267,7 @@ def test_installed_pandas_gbq_version_returns_cached():
 
 def test_installed_pandas_gbq_version_returns_parsed_version():
     import sys
+
     mock_pandas_gbq = mock.Mock()
     mock_pandas_gbq.__version__ = "1.2.3"
     versions = _versions_helpers.PandasGBQVersions()
@@ -268,6 +281,7 @@ def test_installed_pandas_gbq_version_returns_parsed_version():
 
 def test_installed_pandas_gbq_version_falls_back_on_import_error():
     import sys
+
     versions = _versions_helpers.PandasGBQVersions()
     with mock.patch.dict(sys.modules, {"pandas_gbq": None}):
         version = versions.installed_version
@@ -279,6 +293,7 @@ def test_installed_pandas_gbq_version_falls_back_on_import_error():
 
 def test_installed_pandas_gbq_version_falls_back_on_other_error():
     import sys
+
     # Simulate a corrupted package raising an error on import/property access
     class CorruptPandasGBQ:
         @property
@@ -302,6 +317,7 @@ def test_pandas_gbq_delegation_api_version_returns_cached():
 
 def test_pandas_gbq_delegation_api_version_returns_value():
     import sys
+
     mock_pandas_gbq = mock.Mock()
     mock_pandas_gbq._internal_delegation_api_version = 42
     versions = _versions_helpers.PandasGBQVersions()
@@ -313,6 +329,7 @@ def test_pandas_gbq_delegation_api_version_returns_value():
 
 def test_pandas_gbq_delegation_api_version_falls_back_on_import_error():
     import sys
+
     versions = _versions_helpers.PandasGBQVersions()
     with mock.patch.dict(sys.modules, {"pandas_gbq": None}):
         version = versions.delegation_api_version
@@ -322,6 +339,7 @@ def test_pandas_gbq_delegation_api_version_falls_back_on_import_error():
 
 def test_pandas_gbq_delegation_api_version_falls_back_on_other_error():
     import sys
+
     class CorruptPandasGBQ:
         @property
         def _internal_delegation_api_version(self):

packages/google-cloud-bigquery/tests/unit/test_magics.py

@@ -147,6 +147,8 @@ def test_context_with_default_credentials():
     """When Application Default Credentials are set, the context credentials
     will be created the first time it is called
     """
+    magics.context._credentials = None
+    magics.context._project = None
     assert magics.context._credentials is None
     assert magics.context._project is None
 
@@ -164,6 +166,16 @@ def test_context_with_default_credentials():
     assert default_mock.call_count == 2
 
 
+def test_context_fallback_when_bigquery_magics_none():
+    ctx = magics.Context()
+    credentials_mock = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+    with mock.patch("google.auth.default", return_value=(credentials_mock, "proj-123")):
+        assert ctx.credentials is credentials_mock
+        assert ctx.project == "proj-123"
+
+
 @pytest.mark.usefixtures("ipython_interactive")
 @pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
 def test_context_with_default_connection(monkeypatch):