chore: Move remaining test utils to bigframes.testing module (#1810)

TrevorBergeron · sycai · web-flow · commit b3db51974442 · 2025-06-10T20:34:00.000-07:00
Co-authored-by: Shenyang Cai &lt;sycai@users.noreply.github.com&gt;
diff --git a/bigframes/testing/compiler_session.py b/bigframes/testing/compiler_session.py
diff --git a/bigframes/testing/engine_utils.py b/bigframes/testing/engine_utils.py
diff --git a/bigframes/testing/utils.py b/bigframes/testing/utils.py
diff --git a/pytest.ini b/pytest.ini
@@ -2,3 +2,4 @@
 doctest_optionflags = NORMALIZE_WHITESPACE
 filterwarnings =
     ignore::pandas.errors.SettingWithCopyWarning
+addopts = "--import-mode=importlib"
diff --git a/scripts/__init__.py b/scripts/__init__.py
diff --git a/scripts/conftest.py b/scripts/conftest.py
@@ -0,0 +1,8 @@
+from pathlib import Path
+import sys
+
+# inserts scripts into path so that tests can import
+project_root = Path(__file__).parent.parent
+scripts_dir = project_root / "scripts"
+
+sys.path.insert(0, str(scripts_dir))
diff --git a/scripts/test_publish_api_coverage.py b/scripts/test_publish_api_coverage.py
@@ -15,16 +15,15 @@
 import sys
 
 import pandas
+from publish_api_coverage import build_api_coverage_table
 import pytest
 
-from . import publish_api_coverage
-
 pytest.importorskip("sklearn")
 
 
 @pytest.fixture
 def api_coverage_df():
-    return publish_api_coverage.build_api_coverage_table("my_bf_ver", "my_release_ver")
+    return build_api_coverage_table("my_bf_ver", "my_release_ver")
 
 
 @pytest.mark.skipif(
diff --git a/tests/system/conftest.py b/tests/system/conftest.py
@@ -41,7 +41,7 @@
 import bigframes.dataframe
 import bigframes.pandas as bpd
 import bigframes.series
-import tests.system.utils
+import bigframes.testing.utils
 
 # Use this to control the number of cloud functions being deleted in a single
 # test session. This should help soften the spike of the number of mutations per
@@ -615,7 +615,7 @@ def scalars_pandas_df_default_index() -> pd.DataFrame:
         DATA_DIR / "scalars.jsonl",
         lines=True,
     )
-    tests.system.utils.convert_pandas_dtypes(df, bytes_col=True)
+    bigframes.testing.utils.convert_pandas_dtypes(df, bytes_col=True)
 
     df = df.set_index("rowindex", drop=False)
     df.index.name = None
@@ -1422,12 +1422,12 @@ def use_fast_query_path():
 @pytest.fixture(scope="session", autouse=True)
 def cleanup_cloud_functions(session, cloudfunctions_client, dataset_id_permanent):
     """Clean up stale cloud functions."""
-    permanent_endpoints = tests.system.utils.get_remote_function_endpoints(
+    permanent_endpoints = bigframes.testing.utils.get_remote_function_endpoints(
         session.bqclient, dataset_id_permanent
     )
     delete_count = 0
     try:
-        for cloud_function in tests.system.utils.get_cloud_functions(
+        for cloud_function in bigframes.testing.utils.get_cloud_functions(
             cloudfunctions_client,
             session.bqclient.project,
             session.bqclient.location,
@@ -1447,7 +1447,7 @@ def cleanup_cloud_functions(session, cloudfunctions_client, dataset_id_permanent
 
             # Go ahead and delete
             try:
-                tests.system.utils.delete_cloud_function(
+                bigframes.testing.utils.delete_cloud_function(
                     cloudfunctions_client, cloud_function.name
                 )
                 delete_count += 1
diff --git a/tests/system/large/functions/test_managed_function.py b/tests/system/large/functions/test_managed_function.py
@@ -21,7 +21,7 @@
 import bigframes
 import bigframes.exceptions as bfe
 import bigframes.pandas as bpd
-from tests.system.utils import cleanup_function_assets
+from bigframes.testing.utils import cleanup_function_assets
 
 prefixer = test_utils.prefixer.Prefixer("bigframes", "")
 
diff --git a/tests/system/large/functions/test_remote_function.py b/tests/system/large/functions/test_remote_function.py
@@ -36,7 +36,7 @@
 import bigframes.functions._utils as bff_utils
 import bigframes.pandas as bpd
 import bigframes.series
-from tests.system.utils import (
+from bigframes.testing.utils import (
     assert_pandas_df_equal,
     cleanup_function_assets,
     delete_cloud_function,
diff --git a/tests/system/large/ml/test_cluster.py b/tests/system/large/ml/test_cluster.py
@@ -15,7 +15,7 @@
 import pandas as pd
 
 from bigframes.ml import cluster
-from tests.system import utils
+from bigframes.testing import utils
 
 
 def test_cluster_configure_fit_score_predict(
diff --git a/tests/system/large/ml/test_compose.py b/tests/system/large/ml/test_compose.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from bigframes.ml import compose, preprocessing
-from tests.system import utils
+from bigframes.testing import utils
 
 
 def test_columntransformer_standalone_fit_and_transform(
diff --git a/tests/system/large/ml/test_core.py b/tests/system/large/ml/test_core.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from bigframes.ml import globals
-from tests.system import utils
+from bigframes.testing import utils
 
 
 def test_bqml_e2e(session, dataset_id, penguins_df_default_index, new_penguins_df):
diff --git a/tests/system/large/ml/test_decomposition.py b/tests/system/large/ml/test_decomposition.py
@@ -16,7 +16,7 @@
 import pandas.testing
 
 from bigframes.ml import decomposition
-from tests.system import utils
+from bigframes.testing import utils
 
 
 def test_decomposition_configure_fit_score_predict(
diff --git a/tests/system/large/ml/test_ensemble.py b/tests/system/large/ml/test_ensemble.py
@@ -15,7 +15,7 @@
 import pytest
 
 import bigframes.ml.ensemble
-from tests.system import utils
+from bigframes.testing import utils
 
 
 @pytest.mark.flaky(retries=2)
diff --git a/tests/system/large/ml/test_forecasting.py b/tests/system/large/ml/test_forecasting.py
@@ -15,7 +15,7 @@
 import pytest
 
 from bigframes.ml import forecasting
-from tests.system import utils
+from bigframes.testing import utils
 
 ARIMA_EVALUATE_OUTPUT_COL = [
     "non_seasonal_p",
diff --git a/tests/system/large/ml/test_linear_model.py b/tests/system/large/ml/test_linear_model.py
@@ -16,7 +16,7 @@
 
 from bigframes.ml import model_selection
 import bigframes.ml.linear_model
-from tests.system import utils
+from bigframes.testing import utils
 
 
 def test_linear_regression_configure_fit_score(penguins_df_default_index, dataset_id):
diff --git a/tests/system/large/ml/test_model_selection.py b/tests/system/large/ml/test_model_selection.py
@@ -15,7 +15,7 @@
 import pytest
 
 from bigframes.ml import linear_model, model_selection
-from tests.system import utils
+from bigframes.testing import utils
 
 
 @pytest.mark.parametrize(
diff --git a/tests/system/large/ml/test_pipeline.py b/tests/system/large/ml/test_pipeline.py
@@ -25,7 +25,7 @@
     pipeline,
     preprocessing,
 )
-from tests.system import utils
+from bigframes.testing import utils
 
 
 def test_pipeline_linear_regression_fit_score_predict(
diff --git a/tests/system/load/test_llm.py b/tests/system/load/test_llm.py
@@ -16,7 +16,7 @@
 import pytest
 
 from bigframes.ml import llm
-from tests.system import utils
+from bigframes.testing import utils
 
 
 @pytest.fixture(scope="session")
diff --git a/tests/system/small/bigquery/test_vector_search.py b/tests/system/small/bigquery/test_vector_search.py
@@ -23,7 +23,7 @@
 
 import bigframes.bigquery as bbq
 import bigframes.pandas as bpd
-from tests.system.utils import assert_pandas_df_equal
+from bigframes.testing.utils import assert_pandas_df_equal
 
 # Need at least 5,000 rows to create a vector index.
 VECTOR_DF = pd.DataFrame(
diff --git a/tests/system/small/engines/test_read_local.py b/tests/system/small/engines/test_read_local.py
@@ -17,7 +17,7 @@
 import bigframes
 from bigframes.core import identifiers, local_data, nodes
 from bigframes.session import polars_executor
-from tests.system.small.engines.engine_utils import assert_equivalence_execution
+from bigframes.testing.engine_utils import assert_equivalence_execution
 
 pytest.importorskip("polars")
 
diff --git a/tests/system/small/engines/test_selection.py b/tests/system/small/engines/test_selection.py
@@ -16,7 +16,7 @@
 
 from bigframes.core import array_value, expression, identifiers, nodes
 from bigframes.session import polars_executor
-from tests.system.small.engines.engine_utils import assert_equivalence_execution
+from bigframes.testing.engine_utils import assert_equivalence_execution
 
 pytest.importorskip("polars")
 
diff --git a/tests/system/small/functions/test_remote_function.py b/tests/system/small/functions/test_remote_function.py
@@ -31,7 +31,7 @@
 from bigframes.functions import _utils as bff_utils
 from bigframes.functions import function as bff
 import bigframes.session._io.bigquery
-from tests.system.utils import assert_pandas_df_equal, get_function_name
+from bigframes.testing.utils import assert_pandas_df_equal, get_function_name
 
 _prefixer = test_utils.prefixer.Prefixer("bigframes", "")
 
diff --git a/tests/system/small/geopandas/test_geoseries.py b/tests/system/small/geopandas/test_geoseries.py
@@ -31,7 +31,7 @@
 import bigframes.geopandas
 import bigframes.pandas
 import bigframes.series
-from tests.system.utils import assert_series_equal
+from bigframes.testing.utils import assert_series_equal
 
 
 @pytest.fixture(scope="session")
diff --git a/tests/system/small/ml/test_cluster.py b/tests/system/small/ml/test_cluster.py
@@ -16,7 +16,7 @@
 
 from bigframes.ml import cluster
 import bigframes.pandas as bpd
-from tests.system.utils import assert_pandas_df_equal
+from bigframes.testing.utils import assert_pandas_df_equal
 
 _PD_NEW_PENGUINS = pd.DataFrame.from_dict(
     {
diff --git a/tests/system/small/ml/test_core.py b/tests/system/small/ml/test_core.py
@@ -23,7 +23,7 @@
 import bigframes
 import bigframes.features
 from bigframes.ml import core
-from tests.system import utils
+from bigframes.testing import utils
 
 
 def test_model_eval(
diff --git a/tests/system/small/ml/test_decomposition.py b/tests/system/small/ml/test_decomposition.py
@@ -16,7 +16,7 @@
 
 from bigframes.ml import decomposition
 import bigframes.pandas as bpd
-import tests.system.utils
+import bigframes.testing.utils
 
 
 def test_pca_predict(
@@ -33,7 +33,7 @@ def test_pca_predict(
         index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"),
     )
 
-    tests.system.utils.assert_pandas_df_equal_pca(
+    bigframes.testing.utils.assert_pandas_df_equal_pca(
         predictions, expected, check_exact=False, rtol=0.1
     )
 
@@ -161,7 +161,7 @@ def test_pca_components_(penguins_pca_model: decomposition.PCA):
         .reset_index(drop=True)
     )
 
-    tests.system.utils.assert_pandas_df_equal_pca_components(
+    bigframes.testing.utils.assert_pandas_df_equal_pca_components(
         result,
         expected,
         check_exact=False,
@@ -180,7 +180,7 @@ def test_pca_explained_variance_(penguins_pca_model: decomposition.PCA):
             "explained_variance": [3.278657, 1.270829, 1.125354],
         },
     )
-    tests.system.utils.assert_pandas_df_equal(
+    bigframes.testing.utils.assert_pandas_df_equal(
         result,
         expected,
         check_exact=False,
@@ -200,7 +200,7 @@ def test_pca_explained_variance_ratio_(penguins_pca_model: decomposition.PCA):
             "explained_variance_ratio": [0.469357, 0.181926, 0.1611],
         },
     )
-    tests.system.utils.assert_pandas_df_equal(
+    bigframes.testing.utils.assert_pandas_df_equal(
         result,
         expected,
         check_exact=False,
diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py
@@ -22,7 +22,7 @@
 from bigframes import exceptions
 from bigframes.ml import core, llm
 import bigframes.pandas as bpd
-from tests.system import utils
+from bigframes.testing import utils
 
 
 @pytest.mark.parametrize(
diff --git a/tests/system/small/ml/test_multimodal_llm.py b/tests/system/small/ml/test_multimodal_llm.py
@@ -18,7 +18,7 @@
 
 from bigframes.ml import llm
 import bigframes.pandas as bpd
-from tests.system import utils
+from bigframes.testing import utils
 
 
 @pytest.mark.flaky(retries=2)
diff --git a/tests/system/small/ml/test_preprocessing.py b/tests/system/small/ml/test_preprocessing.py
@@ -19,7 +19,7 @@
 
 import bigframes.features
 from bigframes.ml import preprocessing
-from tests.system import utils
+from bigframes.testing import utils
 
 ONE_HOT_ENCODED_DTYPE = (
     pd.ArrowDtype(pa.list_(pa.struct([("index", pa.int64()), ("value", pa.float64())])))
diff --git a/tests/system/small/operations/test_datetimes.py b/tests/system/small/operations/test_datetimes.py
@@ -20,7 +20,7 @@
 import pytest
 
 import bigframes.series
-from tests.system.utils import assert_series_equal
+from bigframes.testing.utils import assert_series_equal
 
 DATETIME_COL_NAMES = [("datetime_col",), ("timestamp_col",)]
 DATE_COLUMNS = [
diff --git a/tests/system/small/operations/test_lists.py b/tests/system/small/operations/test_lists.py
@@ -18,7 +18,7 @@
 import pyarrow as pa
 import pytest
 
-from ...utils import assert_series_equal
+from bigframes.testing.utils import assert_series_equal
 
 
 @pytest.mark.parametrize(
diff --git a/tests/system/small/operations/test_strings.py b/tests/system/small/operations/test_strings.py
@@ -20,8 +20,7 @@
 
 import bigframes.dtypes as dtypes
 import bigframes.pandas as bpd
-
-from ...utils import assert_series_equal
+from bigframes.testing.utils import assert_series_equal
 
 
 def test_find(scalars_dfs):
diff --git a/tests/system/small/regression/test_issue355_merge_after_filter.py b/tests/system/small/regression/test_issue355_merge_after_filter.py
@@ -15,7 +15,7 @@
 import pandas as pd
 import pytest
 
-from tests.system.utils import assert_pandas_df_equal
+from bigframes.testing.utils import assert_pandas_df_equal
 
 
 @pytest.mark.parametrize(
diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py
@@ -33,7 +33,7 @@
 import bigframes.dtypes as dtypes
 import bigframes.pandas as bpd
 import bigframes.series as series
-from tests.system.utils import (
+from bigframes.testing.utils import (
     assert_dfs_equivalent,
     assert_pandas_df_equal,
     assert_series_equal,
diff --git a/tests/system/small/test_dataframe_io.py b/tests/system/small/test_dataframe_io.py
@@ -23,7 +23,7 @@
 import pytest
 
 import bigframes.dtypes as dtypes
-from tests.system import utils
+from bigframes.testing import utils
 
 try:
     import pandas_gbq  # type: ignore
diff --git a/tests/system/small/test_encryption.py b/tests/system/small/test_encryption.py
@@ -21,7 +21,7 @@
 
 import bigframes
 import bigframes.ml.linear_model
-from tests.system import utils
+from bigframes.testing import utils
 
 
 @pytest.fixture(scope="module")
diff --git a/tests/system/small/test_groupby.py b/tests/system/small/test_groupby.py
diff --git a/tests/system/small/test_index.py b/tests/system/small/test_index.py
diff --git a/tests/system/small/test_large_local_data.py b/tests/system/small/test_large_local_data.py
diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py
diff --git a/tests/system/small/test_pandas.py b/tests/system/small/test_pandas.py
diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py
diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py
diff --git a/tests/system/small/test_unordered.py b/tests/system/small/test_unordered.py
diff --git a/tests/unit/core/compile/sqlglot/conftest.py b/tests/unit/core/compile/sqlglot/conftest.py
diff --git a/tests/unit/test_dataframe_polars.py b/tests/unit/test_dataframe_polars.py
