use positional arg for draws like in actual sample

Author: williambdean

pymc/testing.py

@@ -986,7 +986,7 @@ def assert_no_rvs(vars: Sequence[Variable]) -> None:
         raise AssertionError(f"RV found in graph: {rvs}")
 
 
-def mock_sample(*args, **kwargs):
+def mock_sample(draws: int = 10, **kwargs):
     """Mock the pm.sample function by returning prior predictive samples as posterior.
 
     Useful for testing models that use pm.sample without running MCMC sampling.
@@ -1015,7 +1015,7 @@ def mock_pymc_sample():
     """
     random_seed = kwargs.get("random_seed", None)
     model = kwargs.get("model", None)
-    draws = kwargs.get("draws", 10)
+    draws = kwargs.get("draws", draws)
     n_chains = kwargs.get("chains", 1)
     idata: InferenceData = pm.sample_prior_predictive(
         model=model,

tests/test_testing.py

@@ -35,13 +35,24 @@ def test_domain(values, edges, expectation):
         Domain(values, edges=edges)
 
 
-def test_mock_sample() -> None:
+@pytest.mark.parametrize(
+    "args, kwargs, expected_draws",
+    [
+        pytest.param((), {}, 10, id="default"),
+        pytest.param((100,), {}, 100, id="positional-draws"),
+        pytest.param((), {"draws": 100}, 100, id="keyword-draws"),
+    ],
+)
+def test_mock_sample(args, kwargs, expected_draws) -> None:
     _, model, _ = simple_normal(bounded_prior=True)
 
-    idata = mock_sample(model=model)
+    with model:
+        idata = mock_sample(*args, **kwargs)
 
     assert "posterior" in idata
     assert "observed_data" in idata
     assert "prior" not in idata
     assert "posterior_predictive" not in idata
     assert "sample_stats" not in idata
+
+    assert idata.posterior.sizes == {"chain": 1, "draw": expected_draws}