tests

Author: GarrettWu

GEMINI.md

@@ -25,6 +25,7 @@
     explain_predict,
     generate_embedding,
     generate_text,
+    get_insights,
     global_explain,
     predict,
     transform,
@@ -39,4 +40,5 @@
     "transform",
     "generate_text",
     "generate_embedding",
+    "get_insights",
 ]

bigframes/bigquery/ml.py

@@ -64,6 +64,32 @@ def test_generate_embedding_with_options(embedding_model):
     assert len(embedding[0]) == 256
 
 
+def test_get_insights(dataset_id):
+    df = bpd.DataFrame(
+        {
+            "dim1": ["a", "a", "b", "b", "a", "a", "b", "b"],
+            "dim2": ["x", "y", "x", "y", "x", "y", "x", "y"],
+            "metric": [10, 20, 30, 40, 12, 25, 35, 45],
+            "is_test": [False, False, False, False, True, True, True, True],
+        }
+    )
+    model_name = f"{dataset_id}.contribution_analysis_model"
+
+    ml.create_model(
+        model_name=model_name,
+        options={
+            "model_type": "CONTRIBUTION_ANALYSIS",
+            "contribution_metric": "SUM(metric)",
+            "is_test_col": "is_test",
+        },
+        training_data=df,
+    )
+
+    result = ml.get_insights(model_name)
+    assert len(result) > 0
+    assert "contributors" in result.columns
+
+
 def test_create_model_linear_regression(dataset_id):
     df = bpd.DataFrame({"x": [1, 2, 3], "y": [2, 4, 6]})
     model_name = f"{dataset_id}.linear_regression_model"

tests/system/large/bigquery/test_ml.py

@@ -177,6 +177,15 @@ def test_generate_text_with_pandas_dataframe(read_pandas_mock, read_gbq_query_mo
     assert "'TYPE' AS request_type" in generated_sql
 
 
+@mock.patch("bigframes.pandas.read_gbq_query")
+def test_get_insights(read_gbq_query_mock):
+    ml_ops.get_insights(MODEL_SERIES)
+    read_gbq_query_mock.assert_called_once()
+    generated_sql = read_gbq_query_mock.call_args[0][0]
+    assert "ML.GET_INSIGHTS" in generated_sql
+    assert f"MODEL `{MODEL_NAME}`" in generated_sql
+
+
 @mock.patch("bigframes.pandas.read_gbq_query")
 @mock.patch("bigframes.pandas.read_pandas")
 def test_generate_embedding_with_pandas_dataframe(

tests/unit/bigquery/test_ml.py

@@ -203,6 +203,13 @@ def test_generate_text_model_with_options(snapshot):
     snapshot.assert_match(sql, "generate_text_model_with_options.sql")
 
 
+def test_get_insights_model_basic(snapshot):
+    sql = bigframes.core.sql.ml.get_insights(
+        model_name="my_project.my_dataset.my_model",
+    )
+    snapshot.assert_match(sql, "get_insights_model_basic.sql")
+
+
 def test_generate_embedding_model_basic(snapshot):
     sql = bigframes.core.sql.ml.generate_embedding(
         model_name="my_project.my_dataset.my_model",