chore: move multimodal tests from PROD to TEST (#1578)

* add a new session and connect to test env

* I need to revert code change

* revert files for unncessary change

* fix more tests

* all test are clean

* fix more tests

* merge in the change

* revert the order changes, move single-modality tests back to PROD since they are in GA

* refactor multimodal tests out

* change the project name

* revert change

* remove unused bq client

Author: shuoweil

tests/system/conftest.py

@@ -184,6 +184,20 @@ def session_tokyo(tokyo_location: str) -> Generator[bigframes.Session, None, Non
     session.close()  # close generated session at cleanup type
 
 
+@pytest.fixture(scope="session")
+def test_session() -> Generator[bigframes.Session, None, None]:
+    context = bigframes.BigQueryOptions(
+        client_endpoints_override={
+            "bqclient": "https://test-bigquery.sandbox.google.com",
+            "bqconnectionclient": "test-bigqueryconnection.sandbox.googleapis.com",
+            "bqstoragereadclient": "test-bigquerystorage-grpc.sandbox.googleapis.com",
+        },
+    )
+    session = bigframes.Session(context=context)
+    yield session
+    session.close()
+
+
 @pytest.fixture(scope="session")
 def bq_connection_name() -> str:
     return "bigframes-rf-conn"
@@ -910,8 +924,8 @@ def llm_text_pandas_df():
 
 
 @pytest.fixture(scope="session")
-def llm_text_df(session, llm_text_pandas_df):
-    return session.read_pandas(llm_text_pandas_df)
+def llm_text_df(test_session, llm_text_pandas_df):
+    return test_session.read_pandas(llm_text_pandas_df)
 
 
 @pytest.fixture(scope="session")
@@ -1249,10 +1263,10 @@ def penguins_randomforest_classifier_model_name(
 
 @pytest.fixture(scope="session")
 def llm_fine_tune_df_default_index(
-    session: bigframes.Session,
+    test_session: bigframes.Session,
 ) -> bigframes.dataframe.DataFrame:
     training_table_name = "llm_tuning.emotion_classification_train"
-    df = session.read_gbq(training_table_name).dropna().head(30)
+    df = test_session.read_gbq(training_table_name).dropna().head(30)
     prefix = "Please do sentiment analysis on the following text and only output a number from 0 to 5 where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: "
     df["prompt"] = prefix + df["text"]
     df["label"] = df["label"].astype("string")
@@ -1484,13 +1498,14 @@ def images_uris() -> list[str]:
 
 @pytest.fixture(scope="session")
 def images_mm_df(
-    images_gcs_path, session: bigframes.Session, bq_connection: str
+    images_uris, test_session: bigframes.Session, bq_connection: str
 ) -> bpd.DataFrame:
     bigframes.options.experiments.blob = True
 
-    return session.from_glob_path(
-        images_gcs_path, name="blob_col", connection=bq_connection
+    blob_series = bpd.Series(images_uris, session=test_session).str.to_blob(
+        connection=bq_connection
     )
+    return blob_series.rename("blob_col").to_frame()
 
 
 @pytest.fixture()
@@ -1509,8 +1524,10 @@ def pdf_gcs_path() -> str:
 
 @pytest.fixture(scope="session")
 def pdf_mm_df(
-    pdf_gcs_path, session: bigframes.Session, bq_connection: str
+    pdf_gcs_path, test_session: bigframes.Session, bq_connection: str
 ) -> bpd.DataFrame:
     bigframes.options.experiments.blob = True
 
-    return session.from_glob_path(pdf_gcs_path, name="pdf", connection=bq_connection)
+    return test_session.from_glob_path(
+        pdf_gcs_path, name="pdf", connection=bq_connection
+    )

tests/system/large/blob/test_function.py

@@ -55,11 +55,11 @@ def test_blob_image_blur_to_series(
     images_mm_df: bpd.DataFrame,
     bq_connection: str,
     images_output_uris: list[str],
-    session: bigframes.Session,
+    test_session: bigframes.Session,
 ):
     bigframes.options.experiments.blob = True
 
-    series = bpd.Series(images_output_uris, session=session).str.to_blob(
+    series = bpd.Series(images_output_uris, session=test_session).str.to_blob(
         connection=bq_connection
     )
 
@@ -129,11 +129,11 @@ def test_blob_image_resize_to_series(
     images_mm_df: bpd.DataFrame,
     bq_connection: str,
     images_output_uris: list[str],
-    session: bigframes.Session,
+    test_session: bigframes.Session,
 ):
     bigframes.options.experiments.blob = True
 
-    series = bpd.Series(images_output_uris, session=session).str.to_blob(
+    series = bpd.Series(images_output_uris, session=test_session).str.to_blob(
         connection=bq_connection
     )
 
@@ -205,11 +205,11 @@ def test_blob_image_normalize_to_series(
     images_mm_df: bpd.DataFrame,
     bq_connection: str,
     images_output_uris: list[str],
-    session: bigframes.Session,
+    test_session: bigframes.Session,
 ):
     bigframes.options.experiments.blob = True
 
-    series = bpd.Series(images_output_uris, session=session).str.to_blob(
+    series = bpd.Series(images_output_uris, session=test_session).str.to_blob(
         connection=bq_connection
     )
 

tests/system/small/blob/test_io.py

@@ -19,11 +19,11 @@
 
 
 def test_blob_create_from_uri_str(
-    bq_connection: str, session: bigframes.Session, images_uris
+    bq_connection: str, test_session: bigframes.Session, images_uris
 ):
     bigframes.options.experiments.blob = True
 
-    uri_series = bpd.Series(images_uris, session=session)
+    uri_series = bpd.Series(images_uris, session=test_session)
     blob_series = uri_series.str.to_blob(connection=bq_connection)
 
     pd_blob_df = blob_series.struct.explode().to_pandas()
@@ -42,14 +42,21 @@ def test_blob_create_from_uri_str(
 
 
 def test_blob_create_from_glob_path(
-    bq_connection: str, session: bigframes.Session, images_gcs_path, images_uris
+    bq_connection: str, test_session: bigframes.Session, images_gcs_path, images_uris
 ):
     bigframes.options.experiments.blob = True
 
-    blob_df = session.from_glob_path(
+    blob_df = test_session.from_glob_path(
         images_gcs_path, connection=bq_connection, name="blob_col"
     )
-    pd_blob_df = blob_df["blob_col"].struct.explode().to_pandas()
+    pd_blob_df = (
+        blob_df["blob_col"]
+        .struct.explode()
+        .to_pandas()
+        .sort_values("uri")
+        .reset_index(drop=True)
+    )
+
     expected_df = pd.DataFrame(
         {
             "uri": images_uris,
@@ -65,14 +72,20 @@ def test_blob_create_from_glob_path(
 
 
 def test_blob_create_read_gbq_object_table(
-    bq_connection: str, session: bigframes.Session, images_gcs_path, images_uris
+    bq_connection: str, test_session: bigframes.Session, images_gcs_path, images_uris
 ):
     bigframes.options.experiments.blob = True
 
-    obj_table = session._create_object_table(images_gcs_path, bq_connection)
+    obj_table = test_session._create_object_table(images_gcs_path, bq_connection)
 
-    blob_df = session.read_gbq_object_table(obj_table, name="blob_col")
-    pd_blob_df = blob_df["blob_col"].struct.explode().to_pandas()
+    blob_df = test_session.read_gbq_object_table(obj_table, name="blob_col")
+    pd_blob_df = (
+        blob_df["blob_col"]
+        .struct.explode()
+        .to_pandas()
+        .sort_values("uri")
+        .reset_index(drop=True)
+    )
     expected_df = pd.DataFrame(
         {
             "uri": images_uris,

tests/system/small/ml/test_llm.py

@@ -106,25 +106,6 @@ def test_create_load_multimodal_embedding_generator_model(
     assert reloaded_model.connection_name == bq_connection
 
 
-@pytest.mark.flaky(retries=2)
-def test_multimodal_embedding_generator_predict_default_params_success(
-    images_mm_df, session, bq_connection
-):
-    bigframes.options.experiments.blob = True
-
-    text_embedding_model = llm.MultimodalEmbeddingGenerator(
-        connection_name=bq_connection, session=session
-    )
-    df = text_embedding_model.predict(images_mm_df).to_pandas()
-    utils.check_pandas_df_schema_and_index(
-        df,
-        columns=utils.ML_MULTIMODAL_GENERATE_EMBEDDING_OUTPUT,
-        index=2,
-        col_exact=False,
-    )
-    assert len(df["ml_generate_embedding_result"][0]) == 1408
-
-
 @pytest.mark.parametrize(
     "model_name",
     (
@@ -241,36 +222,6 @@ def test_gemini_text_generator_multi_cols_predict_success(
     )
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    (
-        "gemini-1.5-pro-001",
-        "gemini-1.5-pro-002",
-        "gemini-1.5-flash-001",
-        "gemini-1.5-flash-002",
-        "gemini-2.0-flash-exp",
-    ),
-)
-@pytest.mark.flaky(retries=2)
-def test_gemini_text_generator_multimodal_input(
-    images_mm_df: bpd.DataFrame, model_name, session, bq_connection
-):
-    bigframes.options.experiments.blob = True
-
-    gemini_text_generator_model = llm.GeminiTextGenerator(
-        model_name=model_name, connection_name=bq_connection, session=session
-    )
-    pd_df = gemini_text_generator_model.predict(
-        images_mm_df, prompt=["Describe", images_mm_df["blob_col"]]
-    ).to_pandas()
-    utils.check_pandas_df_schema_and_index(
-        pd_df,
-        columns=utils.ML_GENERATE_TEXT_OUTPUT + ["blob_col"],
-        index=2,
-        col_exact=False,
-    )
-
-
 # Overrides __eq__ function for comparing as mock.call parameter
 class EqCmpAllDataFrame(bpd.DataFrame):
     def __eq__(self, other):

tests/system/small/ml/test_multimodal_llm.py

@@ -0,0 +1,69 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import pytest
+
+import bigframes
+from bigframes.ml import llm
+import bigframes.pandas as bpd
+from tests.system import utils
+
+
+@pytest.mark.flaky(retries=2)
+def test_multimodal_embedding_generator_predict_default_params_success(
+    images_mm_df, test_session, bq_connection
+):
+    bigframes.options.experiments.blob = True
+
+    text_embedding_model = llm.MultimodalEmbeddingGenerator(
+        connection_name=bq_connection, session=test_session
+    )
+    df = text_embedding_model.predict(images_mm_df).to_pandas()
+    utils.check_pandas_df_schema_and_index(
+        df,
+        columns=utils.ML_MULTIMODAL_GENERATE_EMBEDDING_OUTPUT,
+        index=2,
+        col_exact=False,
+    )
+    assert len(df["ml_generate_embedding_result"][0]) == 1408
+
+
+@pytest.mark.parametrize(
+    "model_name",
+    (
+        "gemini-1.5-pro-001",
+        "gemini-1.5-pro-002",
+        "gemini-1.5-flash-001",
+        "gemini-1.5-flash-002",
+        "gemini-2.0-flash-exp",
+    ),
+)
+@pytest.mark.flaky(retries=2)
+def test_gemini_text_generator_multimodal_input(
+    images_mm_df: bpd.DataFrame, model_name, test_session, bq_connection
+):
+    bigframes.options.experiments.blob = True
+
+    gemini_text_generator_model = llm.GeminiTextGenerator(
+        model_name=model_name, connection_name=bq_connection, session=test_session
+    )
+    pd_df = gemini_text_generator_model.predict(
+        images_mm_df, prompt=["Describe", images_mm_df["blob_col"]]
+    ).to_pandas()
+    utils.check_pandas_df_schema_and_index(
+        pd_df,
+        columns=utils.ML_GENERATE_TEXT_OUTPUT + ["blob_col"],
+        index=2,
+        col_exact=False,
+    )

tests/system/small/test_session.py

@@ -1617,3 +1617,13 @@ def test_read_json_gcs_default_engine(session, scalars_dfs, gcs_folder):
 
     assert df.shape[0] == scalars_df.shape[0]
     pd.testing.assert_series_equal(df.dtypes, scalars_df.dtypes)
+
+
+def test_read_gbq_test(test_session: bigframes.Session):
+    test_project_id = "bigframes-dev"
+    test_dataset_id = "test_env_only"
+    test_table_id = "one_table"
+    table_id = f"{test_project_id}.{test_dataset_id}.{test_table_id}"
+    actual = test_session.read_gbq(table_id).to_pandas()
+
+    assert actual.shape == (1, 1)