Integrate test into existing old one

gyorilab · Oct 9, 2024 · 4eaf6d9 · 4eaf6d9
1 parent f14178e
commit 4eaf6d9
Showing 1 changed file with 7 additions and 17 deletions.
diff --git a/tests/test_ops.py b/tests/test_ops.py
@@ -579,7 +579,13 @@ def test_stratify_excluded_species():
                   cartesian_control=True,
                   concepts_to_stratify=['susceptible_population'])
 
-    assert len(tm.templates) == 5, templates
+    assert len(tm.templates) == 3, templates
+    assert tm.templates[0].subject.name == 'susceptible_population_vax'
+    assert tm.templates[0].outcome.name == 'infected_population'
+    assert tm.templates[0].controller.name == 'infected_population'
+    assert tm.templates[1].subject.name == 'susceptible_population_unvax'
+    assert tm.templates[1].outcome.name == 'infected_population'
+    assert tm.templates[1].controller.name == 'infected_population'
 
 
 def test_stratify_parameter_consistency():
@@ -630,19 +636,3 @@ def test_add_observable_pattern():
     assert 'young' in tm.observables
     obs = tm.observables['young']
     assert obs.expression.args[0] == sympy.Symbol('A_young') + sympy.Symbol('B_young')
-
-
-def test_cointrollers_excluded_cartesian():
-    m = stratify(sir_parameterized,
-                 key='vax',
-                 strata=['vax', 'unvax'],
-                 structure=[],
-                 cartesian_control=True,
-                 concepts_to_stratify=['susceptible_population'])
-    assert len(m.templates) == 3, m.templates
-    assert m.templates[0].subject.name == 'susceptible_population_vax'
-    assert m.templates[0].outcome.name == 'infected_population'
-    assert m.templates[0].controller.name == 'infected_population'
-    assert m.templates[1].subject.name == 'susceptible_population_unvax'
-    assert m.templates[1].outcome.name == 'infected_population'
-    assert m.templates[1].controller.name == 'infected_population'