Add regression test for fix

gyorilab · Apr 24, 2024 · 72003d1 · 72003d1
1 parent 5bd9944
commit 72003d1
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/tests/test_ops.py b/tests/test_ops.py
@@ -548,3 +548,16 @@ def condition(template):
 
     assert len(tm.templates) == 2
     assert tm.templates[1].rate_law.args[0] == sympy.core.numbers.Zero()
+
+
+def test_stratify_excluded_species():
+    from mira.examples.sir import sir_parameterized
+
+    tm = stratify(sir_parameterized,
+                  key='vax',
+                  strata=['vax', 'unvax'],
+                  structure=[],
+                  cartesian_control=True,
+                  concepts_to_stratify=['susceptible_population'])
+
+    assert len(tm.templates) == 5, templates