add test case

googleapis · Dec 5, 2024 · 289a978 · 289a978
1 parent 47977dc
commit 289a978
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/tests/system/small/ml/test_linear_model.py b/tests/system/small/ml/test_linear_model.py
@@ -109,9 +109,12 @@ def test_linear_reg_model_predict(penguins_linear_model, new_penguins_df):
 def test_linear_reg_model_predict_explain(penguins_linear_model, new_penguins_df):
     predictions = penguins_linear_model.predict_explain(new_penguins_df).to_pandas()
     assert predictions.shape == (3, 12)
-    result = predictions[["predicted_body_mass_g"]]
+    result = predictions[["predicted_body_mass_g", "baseline_prediction_value"]]
     expected = pandas.DataFrame(
-        {"predicted_body_mass_g": [4030.1, 3280.8, 3177.9]},
+        {
+            "predicted_body_mass_g": [4030.1, 3280.8, 3177.9],
+            "baseline_prediction_value": [9362.7, 9362.7, 9362.7],
+        },
         dtype="Float64",
         index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"),
     )