diff --git a/mira/metamodel/composition.py b/mira/metamodel/composition.py
index 6f1dae23..e8162de1 100644
--- a/mira/metamodel/composition.py
+++ b/mira/metamodel/composition.py
@@ -157,10 +157,11 @@ def compose_two_models(tm0, tm1):
                                          new_parameters,
                                          new_initials, new_observables)
 
+                    # it's a refinement relation
                     # get the more specific template
-                    # if it's a refinement, we check to see if the outer or
+                    # we check to see if the outer or
                     # inner template is the more refined version
-                    if result == REFINEMENT_OF:
+                    else:
                         if outer_template_is_more_refined:
                             process_template(new_templates, outer_template,
                                              tm0,
diff --git a/notebooks/Model_Composition_Demo.ipynb b/notebooks/Model_Composition_Demo.ipynb
index 97398a58..8d87cdf9 100644
--- a/notebooks/Model_Composition_Demo.ipynb
+++ b/notebooks/Model_Composition_Demo.ipynb
@@ -24,7 +24,7 @@
    "id": "5aca8c9e-abe6-47e4-bf9a-b0a51bbc0688",
    "metadata": {},
    "source": [
-    "## Create concepts to use for synthetic template models"
+    "## Create concepts to use for template models"
    ]
   },
   {
@@ -34,9 +34,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "susceptible = Concept(\n",
-    "    name=\"susceptible_population\", identifiers={\"ido\": \"0000514\"}\n",
-    ")\n",
+    "susceptible = Concept(name=\"susceptible_population\", identifiers={\"ido\": \"0000514\"})\n",
     "hospitalized = Concept(name=\"hospitalized\", identifiers={\"ncit\": \"C25179\"})\n",
     "infected = Concept(name=\"infected_population\", identifiers={\"ido\": \"0000511\"})\n",
     "recovered = Concept(name=\"immune_population\", identifiers={\"ido\": \"0000592\"})\n",
@@ -75,15 +73,35 @@
     "    controller=infected,\n",
     ")\n",
     "\n",
-    "to_quarantine = NaturalConversion(subject=susceptible, outcome=quarantined)\n",
+    "to_quarantine = NaturalConversion(\n",
+    "    subject=susceptible, \n",
+    "    outcome=quarantined\n",
+    ")\n",
     "\n",
-    "from_quarantine = NaturalConversion(subject=quarantined, outcome=susceptible)\n",
+    "from_quarantine = NaturalConversion(\n",
+    "    subject=quarantined, \n",
+    "    outcome=susceptible\n",
+    ")\n",
     "\n",
-    "dying = NaturalConversion(subject=infected, outcome=dead)\n",
+    "dying = NaturalConversion(\n",
+    "    subject=infected, \n",
+    "    outcome=dead\n",
+    ")\n",
     "\n",
-    "hospitalization = NaturalConversion(subject=infected, outcome= hospitalized)\n",
-    "hospitalization_to_recovery = NaturalConversion(subject=hospitalized, outcome=recovered)\n",
-    "hospitalization_to_death = NaturalConversion(subject=hospitalized, outcome=dead)"
+    "hospitalization = NaturalConversion(\n",
+    "    subject=infected, \n",
+    "    outcome=hospitalized\n",
+    ")\n",
+    "\n",
+    "hospitalization_to_recovery = NaturalConversion(\n",
+    "    subject=hospitalized, \n",
+    "    outcome=recovered\n",
+    ")\n",
+    "\n",
+    "hospitalization_to_death = NaturalConversion(\n",
+    "    subject=hospitalized, \n",
+    "    outcome=dead\n",
+    ")"
    ]
   },
   {
@@ -108,10 +126,21 @@
     "    ]\n",
     ")\n",
     "\n",
-    "sir_reinfection = TemplateModel(templates=[infection, recovery, reinfection])\n",
+    "sir_reinfection = TemplateModel(\n",
+    "    templates=[\n",
+    "        infection, \n",
+    "        recovery, \n",
+    "        reinfection\n",
+    "    ]\n",
+    ")\n",
     "\n",
     "sir_quarantined = TemplateModel(\n",
-    "    templates=[infection, to_quarantine, from_quarantine, recovery]\n",
+    "    templates=[\n",
+    "        infection, \n",
+    "        to_quarantine, \n",
+    "        from_quarantine, \n",
+    "        recovery\n",
+    "    ]\n",
     ")\n",
     "\n",
     "sir_dying = TemplateModel(\n",
@@ -127,11 +156,17 @@
     "        infection,\n",
     "        hospitalization,\n",
     "        hospitalization_to_recovery,\n",
-    "        hospitalization_to_death\n",
+    "        hospitalization_to_death,\n",
     "    ]\n",
     ")\n",
     "\n",
-    "model_list = [sir_reinfection, sir_quarantined, sir_dying, sir_hospitalized, sir]"
+    "model_list = [\n",
+    "    sir_reinfection,\n",
+    "    sir_quarantined,\n",
+    "    sir_dying,\n",
+    "    sir_hospitalized,\n",
+    "    sir,\n",
+    "]"
    ]
   },
   {
diff --git a/tests/test_model_composition.py b/tests/test_model_composition.py
index 478d189e..f6301781 100644
--- a/tests/test_model_composition.py
+++ b/tests/test_model_composition.py
@@ -1,17 +1,59 @@
-from mira.metamodel.composition import *
-from mira.sources.amr import model_from_url
-from mira.examples.sir import sir, sir_2_city
+from mira.metamodel.composition import compose_two_models, compose
+from mira.metamodel.template_model import TemplateModel
+from mira.metamodel.templates import *
+from mira.examples.concepts import *
 
+infection = ControlledConversion(
+    subject=susceptible,
+    outcome=infected,
+    controller=infected,
+)
+recovery = NaturalConversion(
+    subject=infected,
+    outcome=recovered,
+)
 
-def test_model_compose():
-    sir_petrinet_url = 'https://raw.githubusercontent.com/DARPA-ASKEM/' \
-                       'Model-Representations/main/petrinet/examples/sir.json'
-    lotka_regnet_url = ('https://raw.githubusercontent.com/DARPA-ASKEM'
-                        '/Model-Representations/main/regnet'
-                        '/examples/lotka_volterra.json')
+reinfection = ControlledConversion(
+    subject=recovered,
+    outcome=infected,
+    controller=infected,
+)
 
-    tm0 = model_from_url(sir_petrinet_url)
-    tm1 = model_from_url(lotka_regnet_url)
+dying = NaturalConversion(
+    subject=infected,
+    outcome=dead
+)
 
-    new_tm = compose([tm0, tm1, sir, sir_2_city])
+sir = TemplateModel(
+    templates=[
+        infection,
+        recovery,
+    ]
+)
 
+sir_reinfection = TemplateModel(
+    templates=[
+        infection,
+        recovery,
+        reinfection
+    ]
+)
+
+sir_dying = TemplateModel(
+    templates=[
+        infection,
+        dying,
+        recovery,
+    ]
+)
+
+
+def test_compose_two_models():
+    composed_model = compose_two_models(sir_reinfection, sir)
+    assert len(composed_model.templates) == 3
+
+
+def test_compose_list():
+    model_list = [sir_reinfection, sir_dying, sir]
+    composed_model = compose(model_list)
+    assert len(composed_model.templates) == 4