added tests to get_plot_data for pymc and skl experiments

Author: lpoug

causalpy/tests/test_integration_pymc_examples.py

@@ -377,7 +377,10 @@ def test_its():
     assert isinstance(ax, np.ndarray) and all(
         isinstance(item, plt.Axes) for item in ax
     ), "ax must be a numpy.ndarray of plt.Axes"
-
+    plot_data = result.get_plot_data()
+    assert isinstance(plot_data, pd.DataFrame), "The returned object is not a pandas DataFrame"
+    expected_columns = ['prediction', 'pred_hdi_lower', 'pred_hdi_upper', 'impact', 'impact_hdi_lower', 'impact_hdi_upper']
+    assert set(expected_columns).issubset(set(plot_data.columns)), f"DataFrame is missing expected columns {expected_columns}"
 
 @pytest.mark.integration
 def test_its_covid():
@@ -414,6 +417,10 @@ def test_its_covid():
     assert isinstance(ax, np.ndarray) and all(
         isinstance(item, plt.Axes) for item in ax
     ), "ax must be a numpy.ndarray of plt.Axes"
+    plot_data = result.get_plot_data()
+    assert isinstance(plot_data, pd.DataFrame), "The returned object is not a pandas DataFrame"
+    expected_columns = ['prediction', 'pred_hdi_lower', 'pred_hdi_upper', 'impact', 'impact_hdi_lower', 'impact_hdi_upper']
+    assert set(expected_columns).issubset(set(plot_data.columns)), f"DataFrame is missing expected columns {expected_columns}"
 
 
 @pytest.mark.integration
@@ -455,7 +462,10 @@ def test_sc():
     assert isinstance(ax, np.ndarray) and all(
         isinstance(item, plt.Axes) for item in ax
     ), "ax must be a numpy.ndarray of plt.Axes"
-
+    plot_data = result.get_plot_data()
+    assert isinstance(plot_data, pd.DataFrame), "The returned object is not a pandas DataFrame"
+    expected_columns = ['prediction', 'pred_hdi_lower', 'pred_hdi_upper', 'impact', 'impact_hdi_lower', 'impact_hdi_upper']
+    assert set(expected_columns).issubset(set(plot_data.columns)), f"DataFrame is missing expected columns {expected_columns}"
 
 @pytest.mark.integration
 def test_sc_brexit():
@@ -501,6 +511,10 @@ def test_sc_brexit():
     assert isinstance(ax, np.ndarray) and all(
         isinstance(item, plt.Axes) for item in ax
     ), "ax must be a numpy.ndarray of plt.Axes"
+    plot_data = result.get_plot_data()
+    assert isinstance(plot_data, pd.DataFrame), "The returned object is not a pandas DataFrame"
+    expected_columns = ['prediction', 'pred_hdi_lower', 'pred_hdi_upper', 'impact', 'impact_hdi_lower', 'impact_hdi_upper']
+    assert set(expected_columns).issubset(set(plot_data.columns)), f"DataFrame is missing expected columns {expected_columns}"
 
 
 @pytest.mark.integration

causalpy/tests/test_integration_skl_examples.py

@@ -111,6 +111,10 @@ def test_its():
     assert isinstance(ax, np.ndarray) and all(
         isinstance(item, plt.Axes) for item in ax
     ), "ax must be a numpy.ndarray of plt.Axes"
+    plot_data = result.get_plot_data()
+    assert isinstance(plot_data, pd.DataFrame), "The returned object is not a pandas DataFrame"
+    expected_columns = ['prediction', 'impact']
+    assert set(expected_columns).issubset(set(plot_data.columns)), f"DataFrame is missing expected columns {expected_columns}"
 
 
 @pytest.mark.integration
@@ -147,6 +151,10 @@ def test_sc():
     assert isinstance(ax, np.ndarray) and all(
         isinstance(item, plt.Axes) for item in ax
     ), "ax must be a numpy.ndarray of plt.Axes"
+    plot_data = result.get_plot_data()
+    assert isinstance(plot_data, pd.DataFrame), "The returned object is not a pandas DataFrame"
+    expected_columns = ['prediction', 'impact']
+    assert set(expected_columns).issubset(set(plot_data.columns)), f"DataFrame is missing expected columns {expected_columns}"
 
 
 @pytest.mark.integration