Skip to content
This repository was archived by the owner on Apr 1, 2026. It is now read-only.

Commit ce686c1

Browse files
GarrettWuGarrettWu
authored
chore: fix google3 import colabsqlviz tests (#2499)
b/490160312
1 parent 077cb2e commit ce686c1

20 files changed

+734
-637
lines changed

‎bigframes/testing/__init__.py‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,5 @@
1717
These modules are provided for testing the BigQuery DataFrames package. The
1818
interface is not considered stable.
1919
"""
20-
from bigframes.testing.utils import (
21-
assert_frame_equal,
22-
assert_index_equal,
23-
assert_series_equal,
24-
)
2520

26-
__all__ = ["assert_frame_equal", "assert_series_equal", "assert_index_equal"]
21+
# Do not import modules contains pytest. (b/490160312)

‎tests/system/small/bigquery/test_datetime.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
import pytest
2020

2121
from bigframes import bigquery
22-
import bigframes.testing
22+
import bigframes.testing.utils
2323

2424
_TIMESTAMP_DTYPE = pd.ArrowDtype(pa.timestamp("us", tz="UTC"))
2525

@@ -41,7 +41,7 @@ def test_unix_seconds(scalars_dfs):
4141
.apply(lambda ts: _to_unix_epoch(ts, "s"))
4242
.astype("Int64")
4343
)
44-
bigframes.testing.assert_series_equal(actual_res, expected_res)
44+
bigframes.testing.utils.assert_series_equal(actual_res, expected_res)
4545

4646

4747
def test_unix_seconds_after_type_casting(int_series):
@@ -54,7 +54,7 @@ def test_unix_seconds_after_type_casting(int_series):
5454
.apply(lambda ts: _to_unix_epoch(ts, "s"))
5555
.astype("Int64")
5656
)
57-
bigframes.testing.assert_series_equal(
57+
bigframes.testing.utils.assert_series_equal(
5858
actual_res, expected_res, check_index_type=False
5959
)
6060

@@ -76,7 +76,7 @@ def test_unix_millis(scalars_dfs):
7676
.apply(lambda ts: _to_unix_epoch(ts, "ms"))
7777
.astype("Int64")
7878
)
79-
bigframes.testing.assert_series_equal(actual_res, expected_res)
79+
bigframes.testing.utils.assert_series_equal(actual_res, expected_res)
8080

8181

8282
def test_unix_millis_after_type_casting(int_series):
@@ -89,7 +89,7 @@ def test_unix_millis_after_type_casting(int_series):
8989
.apply(lambda ts: _to_unix_epoch(ts, "ms"))
9090
.astype("Int64")
9191
)
92-
bigframes.testing.assert_series_equal(
92+
bigframes.testing.utils.assert_series_equal(
9393
actual_res, expected_res, check_index_type=False
9494
)
9595

@@ -111,7 +111,7 @@ def test_unix_micros(scalars_dfs):
111111
.apply(lambda ts: _to_unix_epoch(ts, "us"))
112112
.astype("Int64")
113113
)
114-
bigframes.testing.assert_series_equal(actual_res, expected_res)
114+
bigframes.testing.utils.assert_series_equal(actual_res, expected_res)
115115

116116

117117
def test_unix_micros_after_type_casting(int_series):
@@ -124,7 +124,7 @@ def test_unix_micros_after_type_casting(int_series):
124124
.apply(lambda ts: _to_unix_epoch(ts, "us"))
125125
.astype("Int64")
126126
)
127-
bigframes.testing.assert_series_equal(
127+
bigframes.testing.utils.assert_series_equal(
128128
actual_res, expected_res, check_index_type=False
129129
)
130130

‎tests/system/small/bigquery/test_geo.py‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@
3232
import bigframes.bigquery as bbq
3333
import bigframes.geopandas
3434
import bigframes.session
35-
import bigframes.testing
35+
import bigframes.testing.utils
3636

3737

3838
def test_geo_st_area(session: bigframes.session.Session):
@@ -57,7 +57,7 @@ def test_geo_st_area(session: bigframes.session.Session):
5757
geobf_s_result = bbq.st_area(geobf_s).to_pandas().round(-3)
5858
assert geobf_s_result.iloc[0] >= 1000
5959

60-
bigframes.testing.assert_series_equal(
60+
bigframes.testing.utils.assert_series_equal(
6161
geobf_s_result,
6262
geopd_s_result,
6363
check_dtype=False,
@@ -110,7 +110,7 @@ def test_st_length_various_geometries(session):
110110

111111
# Test default use_spheroid
112112
result_default = st_length(geoseries).to_pandas()
113-
bigframes.testing.assert_series_equal(
113+
bigframes.testing.utils.assert_series_equal(
114114
result_default,
115115
expected_lengths,
116116
rtol=1e-3,
@@ -119,7 +119,7 @@ def test_st_length_various_geometries(session):
119119

120120
# Test explicit use_spheroid=False
121121
result_explicit_false = st_length(geoseries, use_spheroid=False).to_pandas()
122-
bigframes.testing.assert_series_equal(
122+
bigframes.testing.utils.assert_series_equal(
123123
result_explicit_false,
124124
expected_lengths,
125125
rtol=1e-3,
@@ -153,7 +153,7 @@ def test_geo_st_difference_with_geometry_objects(session: bigframes.session.Sess
153153
index=[0, 1, 2],
154154
dtype=geopandas.array.GeometryDtype(),
155155
)
156-
bigframes.testing.assert_series_equal(
156+
bigframes.testing.utils.assert_series_equal(
157157
geobf_s_result,
158158
expected,
159159
check_index_type=False,
@@ -192,7 +192,7 @@ def test_geo_st_difference_with_single_geometry_object(
192192
index=[0, 1, 2],
193193
dtype=geopandas.array.GeometryDtype(),
194194
)
195-
bigframes.testing.assert_series_equal(
195+
bigframes.testing.utils.assert_series_equal(
196196
geobf_s_result,
197197
expected,
198198
check_index_type=False,
@@ -218,7 +218,7 @@ def test_geo_st_difference_with_similar_geometry_objects(
218218
index=[0, 1, 2],
219219
dtype=geopandas.array.GeometryDtype(),
220220
)
221-
bigframes.testing.assert_series_equal(
221+
bigframes.testing.utils.assert_series_equal(
222222
geobf_s_result,
223223
expected,
224224
check_index_type=False,
@@ -274,7 +274,7 @@ def test_geo_st_distance_with_geometry_objects(session: bigframes.session.Sessio
274274
index=[0, 1, 2, 3],
275275
dtype="Float64",
276276
)
277-
bigframes.testing.assert_series_equal(
277+
bigframes.testing.utils.assert_series_equal(
278278
geobf_s_result,
279279
expected,
280280
check_index_type=False,
@@ -321,7 +321,7 @@ def test_geo_st_distance_with_single_geometry_object(
321321
],
322322
dtype="Float64",
323323
)
324-
bigframes.testing.assert_series_equal(
324+
bigframes.testing.utils.assert_series_equal(
325325
geobf_s_result,
326326
expected,
327327
check_index_type=False,
@@ -356,7 +356,7 @@ def test_geo_st_intersection_with_geometry_objects(session: bigframes.session.Se
356356
index=[0, 1, 2],
357357
dtype=geopandas.array.GeometryDtype(),
358358
)
359-
bigframes.testing.assert_series_equal(
359+
bigframes.testing.utils.assert_series_equal(
360360
geobf_s_result,
361361
expected,
362362
check_index_type=False,
@@ -395,7 +395,7 @@ def test_geo_st_intersection_with_single_geometry_object(
395395
index=[0, 1, 2],
396396
dtype=geopandas.array.GeometryDtype(),
397397
)
398-
bigframes.testing.assert_series_equal(
398+
bigframes.testing.utils.assert_series_equal(
399399
geobf_s_result,
400400
expected,
401401
check_index_type=False,
@@ -425,7 +425,7 @@ def test_geo_st_intersection_with_similar_geometry_objects(
425425
index=[0, 1, 2],
426426
dtype=geopandas.array.GeometryDtype(),
427427
)
428-
bigframes.testing.assert_series_equal(
428+
bigframes.testing.utils.assert_series_equal(
429429
geobf_s_result,
430430
expected,
431431
check_index_type=False,
@@ -466,7 +466,7 @@ def test_geo_st_isclosed(session: bigframes.session.Session):
466466
]
467467
expected_series = pd.Series(data=expected_data, dtype="boolean")
468468

469-
bigframes.testing.assert_series_equal(
469+
bigframes.testing.utils.assert_series_equal(
470470
bf_result,
471471
expected_series,
472472
# We default to Int64 (nullable) dtype, but pandas defaults to int64 index.

‎tests/system/small/bigquery/test_sql.py‎

Lines changed: 25 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import bigframes.bigquery as bbq
1818
import bigframes.dtypes as dtypes
1919
import bigframes.pandas as bpd
20-
import bigframes.testing
20+
import bigframes.testing.utils
2121

2222

2323
def test_sql_scalar_for_all_scalar_types(scalars_df_null_index):
@@ -60,7 +60,9 @@ def test_sql_scalar_for_bool_series(scalars_df_index):
6060
result = bbq.sql_scalar("CAST({0} AS INT64)", [series])
6161
expected = series.astype(dtypes.INT_DTYPE)
6262
expected.name = result.name
63-
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
63+
bigframes.testing.utils.assert_series_equal(
64+
result.to_pandas(), expected.to_pandas()
65+
)
6466

6567

6668
@pytest.mark.parametrize(
@@ -84,7 +86,9 @@ def test_sql_scalar_outputs_all_scalar_types(scalars_df_index, column_name):
8486
result = bbq.sql_scalar("{0}", [series])
8587
expected = series
8688
expected.name = result.name
87-
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
89+
bigframes.testing.utils.assert_series_equal(
90+
result.to_pandas(), expected.to_pandas()
91+
)
8892

8993

9094
def test_sql_scalar_for_array_series(repeated_df):
@@ -114,14 +118,18 @@ def test_sql_scalar_for_array_series(repeated_df):
114118
+ repeated_df["numeric_list_col"].list.len()
115119
+ repeated_df["string_list_col"].list.len()
116120
)
117-
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
121+
bigframes.testing.utils.assert_series_equal(
122+
result.to_pandas(), expected.to_pandas()
123+
)
118124

119125

120126
def test_sql_scalar_outputs_array_series(repeated_df):
121127
result = bbq.sql_scalar("{0}", [repeated_df["int_list_col"]])
122128
expected = repeated_df["int_list_col"]
123129
expected.name = result.name
124-
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
130+
bigframes.testing.utils.assert_series_equal(
131+
result.to_pandas(), expected.to_pandas()
132+
)
125133

126134

127135
def test_sql_scalar_for_struct_series(nested_structs_df):
@@ -132,14 +140,18 @@ def test_sql_scalar_for_struct_series(nested_structs_df):
132140
expected = nested_structs_df["person"].struct.field(
133141
"name"
134142
).str.len() + nested_structs_df["person"].struct.field("age")
135-
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
143+
bigframes.testing.utils.assert_series_equal(
144+
result.to_pandas(), expected.to_pandas()
145+
)
136146

137147

138148
def test_sql_scalar_outputs_struct_series(nested_structs_df):
139149
result = bbq.sql_scalar("{0}", [nested_structs_df["person"]])
140150
expected = nested_structs_df["person"]
141151
expected.name = result.name
142-
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
152+
bigframes.testing.utils.assert_series_equal(
153+
result.to_pandas(), expected.to_pandas()
154+
)
143155

144156

145157
def test_sql_scalar_for_json_series(json_df):
@@ -151,11 +163,15 @@ def test_sql_scalar_for_json_series(json_df):
151163
)
152164
expected = bbq.json_value(json_df["json_col"], "$.int_value")
153165
expected.name = result.name
154-
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
166+
bigframes.testing.utils.assert_series_equal(
167+
result.to_pandas(), expected.to_pandas()
168+
)
155169

156170

157171
def test_sql_scalar_outputs_json_series(json_df):
158172
result = bbq.sql_scalar("{0}", [json_df["json_col"]])
159173
expected = json_df["json_col"]
160174
expected.name = result.name
161-
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
175+
bigframes.testing.utils.assert_series_equal(
176+
result.to_pandas(), expected.to_pandas()
177+
)

‎tests/system/small/bigquery/test_struct.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616

1717
import bigframes.bigquery as bbq
1818
import bigframes.series as series
19-
import bigframes.testing
19+
import bigframes.testing.utils
2020

2121

2222
@pytest.mark.parametrize(
@@ -53,7 +53,7 @@ def test_struct_from_dataframe(columns_arg):
5353
srs = series.Series(
5454
columns_arg,
5555
)
56-
bigframes.testing.assert_series_equal(
56+
bigframes.testing.utils.assert_series_equal(
5757
srs.to_pandas(),
5858
bbq.struct(srs.struct.explode()).to_pandas(),
5959
check_index_type=False,

‎tests/system/small/core/test_reshape.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
from bigframes import session
1919
from bigframes.core.reshape import merge
20-
import bigframes.testing
20+
import bigframes.testing.utils
2121

2222

2323
@pytest.mark.parametrize(
@@ -56,7 +56,7 @@ def test_join_with_index(
5656
how=how,
5757
)
5858

59-
bigframes.testing.assert_frame_equal(
59+
bigframes.testing.utils.assert_frame_equal(
6060
bf_result, pd_result, check_dtype=False, check_index_type=False
6161
)
6262

0 commit comments

Comments
 (0)