Update and format test module

Author: jsakv

causalpy/tests/test_pymc_models.py

@@ -22,7 +22,6 @@
 from pymc_extras.prior import Prior
 
 import causalpy as cp
-
 from causalpy.data.simulate_data import generate_hlr_data
 from causalpy.pymc_models import (
     HierarchicalLinearRegression,
@@ -845,7 +844,7 @@ def _prepare_data(
                 "group_idx": group_idx,
                 "coords": coords,
             },
-            "sample_kwargs": sample_kwargs,
+            "sample_kwargs": resolved_sample_kwargs,
             "priors": {
                 "beta_fixed": Prior(
                     "Normal", mu=0, sigma=10, dims=["treated_units", "coeffs"]
@@ -922,7 +921,7 @@ def setup_class(cls) -> None:
             coords=cls.coords,
         )
         with cls.model:
-            cls.idata = pm.sample(**cls.sample_kwargs)
+            cls.idata = pm.sampling.mcmc.sample(**cls.sample_kwargs)
 
     def test_centered_requires_sigma_random_prior(self) -> None:
         priors = {key: self.model.priors[key] for key in ["beta_fixed", "sigma_fixed"]}