diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py index 10a5c59bb..51a33e6dc 100644 --- a/google/cloud/bigquery/_pandas_helpers.py +++ b/google/cloud/bigquery/_pandas_helpers.py @@ -341,7 +341,7 @@ def types_mapper(arrow_data_type): # recognized by coverage, hence the pragma. See Issue: #2132 elif ( range_timestamp_dtype is not None - and arrow_data_type.equals( # pragma: NO COVER + and arrow_data_type.equals( range_timestamp_dtype.pyarrow_dtype ) ): diff --git a/tests/unit/test__pandas_helpers.py b/tests/unit/test__pandas_helpers.py index bc94f5f54..fbc407836 100644 --- a/tests/unit/test__pandas_helpers.py +++ b/tests/unit/test__pandas_helpers.py @@ -51,6 +51,7 @@ from google.cloud.bigquery import _versions_helpers from google.cloud.bigquery import schema from google.cloud.bigquery._pandas_helpers import determine_requested_streams +# from google.cloud.bigquery.dtypes import RangeTIMESTAMPDtype # Commented out as it causes ImportError and test uses a mock pyarrow = _versions_helpers.PYARROW_VERSIONS.try_import() @@ -1294,6 +1295,51 @@ def test_dataframe_to_bq_schema_returns_schema_with_pandas_gbq( assert got is not None +@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`") +@pytest.mark.skipif(pandas is None, reason="Requires `pandas`") # default_types_mapper is pandas-specific +def test_default_types_mapper_range_timestamp(module_under_test): + """Test default_types_mapper with a mock RangeTIMESTAMPDtype.""" + + # Define a duck-typed class that mimics RangeTIMESTAMPDtype's structure + # for the purpose of this test, as the actual class import is problematic. + class MockRangeTSDtype: + def __init__(self): + self.pyarrow_dtype = pyarrow.struct( + [ + ("start", pyarrow.timestamp("us", tz="UTC")), + ("end", pyarrow.timestamp("us", tz="UTC")), + ] + ) + + def __eq__(self, other): + if not isinstance(other, MockRangeTSDtype): + return NotImplemented + return self.pyarrow_dtype == other.pyarrow_dtype + + # PyArrow dtype that the mapper will receive + range_timestamp_pyarrow_dtype = pyarrow.struct( + [ + ("start", pyarrow.timestamp("us", tz="UTC")), + ("end", pyarrow.timestamp("us", tz="UTC")), + ] + ) + + # Instantiate the mock RangeTIMESTAMPDtype + mock_range_ts_dtype = MockRangeTSDtype() + + # Get the mapper function from the module_under_test fixture + mapper = module_under_test.default_types_mapper( + range_timestamp_dtype=mock_range_ts_dtype + ) + + # Call the mapper with the created PyArrow struct + result_dtype = mapper(range_timestamp_pyarrow_dtype) + + # Assert the result + assert isinstance(result_dtype, MockRangeTSDtype) + assert result_dtype == mock_range_ts_dtype + + @pytest.mark.skipif(pandas is None, reason="Requires `pandas`") def test_dataframe_to_bq_schema_w_named_index(module_under_test, monkeypatch): monkeypatch.setattr(module_under_test, "pandas_gbq", None)