test plot functions

Signed-off-by: Nathaniel <NathanielF@users.noreply.github.com>

Author: NathanielF

causalpy/pymc_experiments.py

@@ -1723,6 +1723,8 @@ def make_hists(idata, i, axs):
         axs[2].legend()
         axs[2].set_title("Average Treatment Effect", fontsize=20)
 
+        return fig
+
     def weighted_percentile(self, data, weights, perc):
         """
         perc : percentile in [0-1]!
@@ -1792,3 +1794,4 @@ def plot_balance_ecdf(self, covariate, idata=None, weighting_scheme=None):
         axs[0].set_xlabel("Quantiles")
         axs[1].legend()
         axs[0].legend()
+        return fig

causalpy/tests/test_pymc_experiments.py

@@ -2,6 +2,7 @@
 Unit tests for pymc_experiments.py
 """
 import arviz as az
+import matplotlib as mpl
 import pandas as pd
 
 import causalpy as cp
@@ -77,3 +78,7 @@ def test_inverse_prop():
     assert isinstance(ate_list, list)
     ate_list = result.get_ate(0, result.idata, method="overlap")
     assert isinstance(ate_list, list)
+    fig = result.plot_ATE()
+    assert isinstance(fig, mpl.figure.Figure)
+    fig = result.plot_balance_ecdf("age")
+    assert isinstance(fig, mpl.figure.Figure)