chore: Multimodal tests to target at BQ prod (#1782)

Author: GarrettWu

tests/system/conftest.py

@@ -920,8 +920,8 @@ def llm_text_pandas_df():
 
 
 @pytest.fixture(scope="session")
-def llm_text_df(test_session, llm_text_pandas_df):
-    return test_session.read_pandas(llm_text_pandas_df)
+def llm_text_df(session, llm_text_pandas_df):
+    return session.read_pandas(llm_text_pandas_df)
 
 
 @pytest.fixture(scope="session")
@@ -1494,9 +1494,9 @@ def images_uris() -> list[str]:
 
 @pytest.fixture(scope="session")
 def images_mm_df(
-    images_uris, test_session: bigframes.Session, bq_connection: str
+    images_uris, session: bigframes.Session, bq_connection: str
 ) -> bpd.DataFrame:
-    blob_series = bpd.Series(images_uris, session=test_session).str.to_blob(
+    blob_series = bpd.Series(images_uris, session=session).str.to_blob(
         connection=bq_connection
     )
     return blob_series.rename("blob_col").to_frame()
@@ -1518,8 +1518,6 @@ def pdf_gcs_path() -> str:
 
 @pytest.fixture(scope="session")
 def pdf_mm_df(
-    pdf_gcs_path, test_session: bigframes.Session, bq_connection: str
+    pdf_gcs_path, session: bigframes.Session, bq_connection: str
 ) -> bpd.DataFrame:
-    return test_session.from_glob_path(
-        pdf_gcs_path, name="pdf", connection=bq_connection
-    )
+    return session.from_glob_path(pdf_gcs_path, name="pdf", connection=bq_connection)

tests/system/large/blob/test_function.py

@@ -53,9 +53,9 @@ def images_output_uris(images_output_folder: str) -> list[str]:
 
 def test_blob_exif(
     bq_connection: str,
-    test_session: bigframes.Session,
+    session: bigframes.Session,
 ):
-    exif_image_df = test_session.from_glob_path(
+    exif_image_df = session.from_glob_path(
         "gs://bigframes_blob_test/images_exif/*",
         name="blob_col",
         connection=bq_connection,
@@ -64,7 +64,7 @@ def test_blob_exif(
     actual = exif_image_df["blob_col"].blob.exif(connection=bq_connection)
     expected = bpd.Series(
         ['{"ExifOffset": 47, "Make": "MyCamera"}'],
-        session=test_session,
+        session=session,
         dtype=dtypes.JSON_DTYPE,
     )
     pd.testing.assert_series_equal(
@@ -79,9 +79,9 @@ def test_blob_image_blur_to_series(
     images_mm_df: bpd.DataFrame,
     bq_connection: str,
     images_output_uris: list[str],
-    test_session: bigframes.Session,
+    session: bigframes.Session,
 ):
-    series = bpd.Series(images_output_uris, session=test_session).str.to_blob(
+    series = bpd.Series(images_output_uris, session=session).str.to_blob(
         connection=bq_connection
     )
 
@@ -147,9 +147,9 @@ def test_blob_image_resize_to_series(
     images_mm_df: bpd.DataFrame,
     bq_connection: str,
     images_output_uris: list[str],
-    test_session: bigframes.Session,
+    session: bigframes.Session,
 ):
-    series = bpd.Series(images_output_uris, session=test_session).str.to_blob(
+    series = bpd.Series(images_output_uris, session=session).str.to_blob(
         connection=bq_connection
     )
 
@@ -217,9 +217,9 @@ def test_blob_image_normalize_to_series(
     images_mm_df: bpd.DataFrame,
     bq_connection: str,
     images_output_uris: list[str],
-    test_session: bigframes.Session,
+    session: bigframes.Session,
 ):
-    series = bpd.Series(images_output_uris, session=test_session).str.to_blob(
+    series = bpd.Series(images_output_uris, session=session).str.to_blob(
         connection=bq_connection
     )
 

tests/system/small/blob/test_io.py

@@ -19,9 +19,9 @@
 
 
 def test_blob_create_from_uri_str(
-    bq_connection: str, test_session: bigframes.Session, images_uris
+    bq_connection: str, session: bigframes.Session, images_uris
 ):
-    uri_series = bpd.Series(images_uris, session=test_session)
+    uri_series = bpd.Series(images_uris, session=session)
     blob_series = uri_series.str.to_blob(connection=bq_connection)
 
     pd_blob_df = blob_series.struct.explode().to_pandas()
@@ -40,9 +40,9 @@ def test_blob_create_from_uri_str(
 
 
 def test_blob_create_from_glob_path(
-    bq_connection: str, test_session: bigframes.Session, images_gcs_path, images_uris
+    bq_connection: str, session: bigframes.Session, images_gcs_path, images_uris
 ):
-    blob_df = test_session.from_glob_path(
+    blob_df = session.from_glob_path(
         images_gcs_path, connection=bq_connection, name="blob_col"
     )
     pd_blob_df = (
@@ -68,11 +68,11 @@ def test_blob_create_from_glob_path(
 
 
 def test_blob_create_read_gbq_object_table(
-    bq_connection: str, test_session: bigframes.Session, images_gcs_path, images_uris
+    bq_connection: str, session: bigframes.Session, images_gcs_path, images_uris
 ):
-    obj_table = test_session._create_object_table(images_gcs_path, bq_connection)
+    obj_table = session._create_object_table(images_gcs_path, bq_connection)
 
-    blob_df = test_session.read_gbq_object_table(obj_table, name="blob_col")
+    blob_df = session.read_gbq_object_table(obj_table, name="blob_col")
     pd_blob_df = (
         blob_df["blob_col"]
         .struct.explode()

tests/system/small/ml/test_multimodal_llm.py

@@ -23,10 +23,10 @@
 
 @pytest.mark.flaky(retries=2)
 def test_multimodal_embedding_generator_predict_default_params_success(
-    images_mm_df, test_session, bq_connection
+    images_mm_df, session, bq_connection
 ):
     text_embedding_model = llm.MultimodalEmbeddingGenerator(
-        connection_name=bq_connection, session=test_session
+        connection_name=bq_connection, session=session
     )
     df = text_embedding_model.predict(images_mm_df).to_pandas()
     utils.check_pandas_df_schema_and_index(
@@ -48,10 +48,10 @@ def test_multimodal_embedding_generator_predict_default_params_success(
 )
 @pytest.mark.flaky(retries=2)
 def test_gemini_text_generator_multimodal_input(
-    images_mm_df: bpd.DataFrame, model_name, test_session, bq_connection
+    images_mm_df: bpd.DataFrame, model_name, session, bq_connection
 ):
     gemini_text_generator_model = llm.GeminiTextGenerator(
-        model_name=model_name, connection_name=bq_connection, session=test_session
+        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"]]
@@ -73,10 +73,10 @@ def test_gemini_text_generator_multimodal_input(
 )
 @pytest.mark.flaky(retries=2)
 def test_gemini_text_generator_multimodal_structured_output(
-    images_mm_df: bpd.DataFrame, model_name, test_session, bq_connection
+    images_mm_df: bpd.DataFrame, model_name, session, bq_connection
 ):
     gemini_text_generator_model = llm.GeminiTextGenerator(
-        model_name=model_name, connection_name=bq_connection, session=test_session
+        model_name=model_name, connection_name=bq_connection, session=session
     )
     output_schema = {
         "bool_output": "bool",

tests/system/small/test_session.py

@@ -1814,16 +1814,6 @@ def test_read_json_gcs_default_engine(session, scalars_dfs, gcs_folder):
     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)
-
-
 @pytest.mark.parametrize(
     ("query_or_table", "index_col", "columns"),
     [