load bug with y name different than output_var (#1468)

* change the name of the series / array

* modifications to tests to use prior_predictive mocking

* link up y in test

* add the prior to the idata

* add different name in xarray test

pymc_marketing/model_builder.py

@@ -632,8 +632,7 @@ def create_fit_data(
         if isinstance(y, np.ndarray):
             y = pd.Series(y, index=X.index, name=self.output_var)
 
-        if y.name is None:
-            y.name = self.output_var
+        y.name = self.output_var
 
         if isinstance(X, pd.DataFrame):
             X = X.to_xarray()

tests/conftest.py

@@ -147,7 +147,7 @@ def mock_sample(*args, **kwargs):
     return idata
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def mock_pymc_sample():
     original_sample = pm.sample
     pm.sample = mock_sample

tests/customer_choice/test_mv_its.py

@@ -13,7 +13,6 @@
 #   limitations under the License.
 
 import re
-import warnings
 
 import numpy as np
 import pandas as pd
@@ -84,38 +83,10 @@ def test_plot_data(saturated_data):
     plt.close()
 
 
-def mock_fit(self, X, y, **kwargs):
-    self.idata.add_groups(
-        {
-            "posterior": self.idata.prior,
-        }
-    )
-
-    combined_data = pd.concat([X, y.rename(self.output_var)], axis=1)
-
-    if "fit_data" in self.idata:
-        del self.idata.fit_data
-
-    with warnings.catch_warnings():
-        warnings.filterwarnings(
-            "ignore",
-            category=UserWarning,
-            message="The group fit_data is not defined in the InferenceData scheme",
-        )
-        self.idata.add_groups(fit_data=combined_data.to_xarray())  # type: ignore
-
-    return self
-
-
 @pytest.fixture(scope="module")
-def fit_model(module_mocker, saturated_data):
+def fit_model(saturated_data, mock_pymc_sample):
     model = MVITS(existing_sales=["competitor", "own"], saturated_market=True)
 
-    module_mocker.patch(
-        "pymc_marketing.customer_choice.mv_its.MVITS.fit",
-        mock_fit,
-    )
-
     model.sample(
         saturated_data.loc[:, ["competitor", "own"]],
         saturated_data["new"],
@@ -151,14 +122,9 @@ def unsaturated_data_good():
 
 
 @pytest.fixture(scope="module")
-def unsaturated_model_bad(module_mocker, unsaturated_data_bad):
+def unsaturated_model_bad(unsaturated_data_bad, mock_pymc_sample):
     model = MVITS(existing_sales=["competitor", "own"], saturated_market=False)
 
-    module_mocker.patch(
-        "pymc_marketing.customer_choice.mv_its.MVITS.fit",
-        mock_fit,
-    )
-
     model.sample(
         unsaturated_data_bad.loc[:, ["competitor", "own"]],
         unsaturated_data_bad["new"],
@@ -169,14 +135,9 @@ def unsaturated_model_bad(module_mocker, unsaturated_data_bad):
 
 
 @pytest.fixture(scope="module")
-def unsaturated_model_good(module_mocker, unsaturated_data_good):
+def unsaturated_model_good(unsaturated_data_good, mock_pymc_sample):
     model = MVITS(existing_sales=["competitor", "own"], saturated_market=False)
 
-    module_mocker.patch(
-        "pymc_marketing.customer_choice.mv_its.MVITS.fit",
-        mock_fit,
-    )
-
     model.sample(
         unsaturated_data_good.loc[:, ["competitor", "own"]],
         unsaturated_data_good["new"],
@@ -220,8 +181,11 @@ def test_save_load(fit_model, saturated_data) -> None:
     assert loaded.saturated_market == fit_model.saturated_market
     assert loaded.X.columns.name is None
     pd.testing.assert_frame_equal(loaded.X, fit_model.X, check_names=False)
-    assert loaded.y.name == fit_model.output_var
-    pd.testing.assert_series_equal(loaded.y.rename("new"), saturated_data["new"])
+    assert loaded.y.name == fit_model.y.name
+    pd.testing.assert_series_equal(
+        loaded.y,
+        saturated_data["new"].rename(fit_model.output_var),
+    )
 
 
 @pytest.mark.parametrize("variable", ["y", "mu"])

tests/mmm/test_base.py

@@ -314,7 +314,12 @@ def test_calling_prior_predictive_before_fit_raises_error(test_mmm, toy_X, toy_y
         test_mmm.prior_predictive
 
 
-def test_calling_fit_result_before_fit_raises_error(test_mmm, toy_X, toy_y):
+def test_calling_fit_result_before_fit_raises_error(
+    test_mmm,
+    toy_X,
+    toy_y,
+    mock_pymc_sample,
+):
     # Arrange
     test_mmm.idata = None
     with pytest.raises(