fix unit tests

Author: sycai

bigframes/bigquery/_operations/ai.py

@@ -26,7 +26,7 @@
 from bigframes import clients, dataframe, dtypes
 from bigframes import pandas as bpd
 from bigframes import series, session
-from bigframes.bigquery._operations import utils as ml_utils
+from bigframes.bigquery._operations import utils as bq_utils
 from bigframes.core import convert
 from bigframes.core.logging import log_adapter
 import bigframes.core.sql.literals
@@ -455,8 +455,8 @@ def generate_embedding(
             for details.
     """
     data = _to_dataframe(data, series_rename="content")
-    model_name, session = ml_utils.get_model_name_and_session(model, data)
-    table_sql = ml_utils.to_sql(data)
+    model_name, session = bq_utils.get_model_name_and_session(model, data)
+    table_sql = bq_utils.to_sql(data)
 
     struct_fields: Dict[str, bigframes.core.sql.literals.STRUCT_VALUES] = {}
     if output_dimensionality is not None:
@@ -564,8 +564,8 @@ def generate_text(
             The generated text.
     """
     data = _to_dataframe(data, series_rename="prompt")
-    model_name, session = ml_utils.get_model_name_and_session(model, data)
-    table_sql = ml_utils.to_sql(data)
+    model_name, session = bq_utils.get_model_name_and_session(model, data)
+    table_sql = bq_utils.to_sql(data)
 
     struct_fields: Dict[
         str,

bigframes/bigquery/_operations/utils.py

@@ -18,6 +18,7 @@
 
 import bigframes
 from bigframes import dataframe
+import bigframes.ml.base
 
 
 def get_model_name_and_session(

tests/unit/bigquery/test_ml.py

@@ -40,31 +40,6 @@ def mock_session():
 MODEL_NAME = "test-project.test-dataset.test-model"
 
 
-def test_get_model_name_and_session_with_pandas_series_model_input():
-    model_name, _ = ml_ops._get_model_name_and_session(MODEL_SERIES)
-    assert model_name == MODEL_NAME
-
-
-def test_get_model_name_and_session_with_pandas_series_model_input_missing_model_reference():
-    model_series = pd.Series({"some_other_key": "value"})
-    with pytest.raises(
-        ValueError, match="modelReference must be present in the pandas Series"
-    ):
-        ml_ops._get_model_name_and_session(model_series)
-
-
-@mock.patch("bigframes.pandas.read_pandas")
-def test_to_sql_with_pandas_dataframe(read_pandas_mock):
-    df = pd.DataFrame({"col1": [1, 2, 3]})
-    read_pandas_mock.return_value._to_sql_query.return_value = (
-        "SELECT * FROM `pandas_df`",
-        [],
-        [],
-    )
-    ml_ops._to_sql(df)
-    read_pandas_mock.assert_called_once()
-
-
 @mock.patch("bigframes.bigquery._operations.ml._get_model_metadata")
 @mock.patch("bigframes.pandas.read_pandas")
 def test_create_model_with_pandas_dataframe(