test: added test for predictions grouping

Author: butterman0

tests/test_model_builder.py

@@ -304,3 +304,35 @@ def test_id():
     ).hexdigest()[:16]
 
     assert model_builder.id == expected_id
+
+@pytest.mark.parametrize("predictions", [True, False])
+def test_predict_respects_predictions_flag(fitted_model_instance, predictions):
+    x_pred = np.random.uniform(0, 1, 100)
+    prediction_data = pd.DataFrame({"input": x_pred})
+    output_var = fitted_model_instance.output_var
+
+    # Snapshot the original posterior_predictive values
+    pp_before = fitted_model_instance.idata.posterior_predictive[output_var].values.copy()
+
+    # Ensure 'predictions' group is not present initially
+    assert "predictions" not in fitted_model_instance.idata.groups()
+
+    # Run prediction with predictions=True or False
+    fitted_model_instance.predict(
+        prediction_data["input"],
+        extend_idata=True,
+        combined=False,
+        predictions=predictions,
+    )
+    
+    pp_after = fitted_model_instance.idata.posterior_predictive[output_var].values
+
+    # Check predictions group presence
+    if predictions:
+        assert "predictions" in fitted_model_instance.idata.groups()
+        # Posterior predictive should remain unchanged
+        np.testing.assert_array_equal(pp_before, pp_after)
+    else:
+        assert "predictions" not in fitted_model_instance.idata.groups()
+        # Posterior predictive should be updated
+        np.testing.assert_array_not_equal(pp_before, pp_after)