diff --git a/tests/test_distribution.py b/tests/test_distribution.py
new file mode 100644
index 0000000..6e35643
--- /dev/null
+++ b/tests/test_distribution.py
@@ -0,0 +1,13 @@
+import epiworldpy as epiworld
+
+def test_distribution():
+    hypothetical = epiworld.ModelSIR(
+        name              = 'hypothetical',
+        prevalence        = 0.01,
+        transmission_rate = 0.1,
+        recovery_rate     = 0.14
+    )
+
+    hypothetical.agents_smallworld(100000, 10, False, 0.01)
+
+    hypothetical.run(100, 223)
diff --git a/tests/test_entity.py b/tests/test_entity.py
new file mode 100644
index 0000000..486a000
--- /dev/null
+++ b/tests/test_entity.py
@@ -0,0 +1,30 @@
+import epiworldpy as epiworld
+
+def dist_factory(start, end):
+    def entity_to_agent_fun(entity, model):
+            agents = model.get_agents()
+
+            for i in range(start, end):
+                entity.add_agent(agents[i], model)
+
+    return entity_to_agent_fun
+
+def test_entity():
+    hypothetical = epiworld.ModelSIRCONN(
+        name              = 'hypothetical',
+        n                 = 10000,
+        prevalence        = 0.01,
+        contact_rate      = 10,
+        transmission_rate = 0.1,
+        recovery_rate     = 0.14
+    )
+
+    entity_1 = epiworld.Entity("Entity 1", dist_factory(0, 3000))
+    entity_2 = epiworld.Entity("Entity 2", dist_factory(3000, 6000))
+    entity_3 = epiworld.Entity("Entity 3", dist_factory(6000, 10000))
+
+    hypothetical.add_entity(entity_1)
+    hypothetical.add_entity(entity_2)
+    hypothetical.add_entity(entity_3)
+
+    hypothetical.run(100, 223)
diff --git a/tests/test_models.py b/tests/test_models.py
index 0ad0a93..7bf487b 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1,75 +1,38 @@
 import epiworldpy as epiworld
 
-# TODO: Decide the right way to validate data; how often is the underlying C++
-# simulation code changed in a way that will change outputs here? What's the
-# tradeoff?
-
-def test_diffnet_simple():
-    # TODO: How do we invoke this model?
-    pass
-
-def test_seir_simple():
-    # TODO: Implement `agents_from_adjlist` or something similar.
-    pass
-
 def test_seirconn_simple():
-    covid19 = epiworld.ModelSEIRCONN(
-        name              = 'covid-19',
+    hypothetical = epiworld.ModelSEIRCONN(
+        name              = 'hypothetical',
         n                 = 10000,
-        prevalence        = .01,
+        prevalence        = 0.01,
         contact_rate      = 2.0,
-        transmission_rate = .1,
+        transmission_rate = 0.1,
         incubation_days   = 7.0,
         recovery_rate     = 0.14
     )
 
-    covid19.run(100, 223)
-
-def test_seird_simple():
-    # TODO: Implement `agents_from_adjlist` or something similar.
-    pass
-
-def test_sir_simple():
-    # TODO: Implement `agents_from_adjlist` or something similar.
-    pass
+    hypothetical.run(100, 223)
 
 def test_sirconn_simple():
-    covid19 = epiworld.ModelSIRCONN(
-        name              = 'covid-19',
+    hypothetical = epiworld.ModelSIRCONN(
+        name              = 'hypothetical',
         n                 = 10000,
-        prevalence        = .01,
+        prevalence        = 0.01,
         contact_rate      = 2.0,
-        transmission_rate = .1,
+        transmission_rate = 0.1,
         recovery_rate     = 0.14
     )
 
-    covid19.run(100, 223)
-
-def test_sird_simple():
-    # TODO: Implement `agents_from_adjlist` or something similar.
-    pass
+    hypothetical.run(100, 223)
 
-def test_sirdconn_simple():
-    covid19 = epiworld.ModelSIRDCONN(
-        name              = 'covid-19',
-        n                 = 10000,
-        prevalence        = .01,
-        contact_rate      = 2.0,
-        transmission_rate = .1,
-        recovery_rate     = 0.14,
-        death_rate        = 0.1,
+def test_smallworld():
+    hypothetical = epiworld.ModelSIR(
+        name              = 'hypothetical',
+        prevalence        = 0.01,
+        transmission_rate = 0.1,
+        recovery_rate     = 0.14
     )
 
-    covid19.run(100, 223)
-
-def test_sis_simple():
-    # TODO: Implement `agents_from_adjlist` or something similar.
-    pass
-
-def test_sisd_simple():
-    # TODO: Implement `agents_from_adjlist` or something similar.
-    pass
+    hypothetical.agents_smallworld(100000, 10, False, 0.01)
 
-def test_surv_simple():
-    # TODO: Implement `agents_from_adjlist` or something similar.
-    pass
+    hypothetical.run(100, 223)
diff --git a/tests/test_saver.py b/tests/test_saver.py
deleted file mode 100644
index a234b47..0000000
--- a/tests/test_saver.py
+++ /dev/null
@@ -1,18 +0,0 @@
-import epiworldpy as epiworld
-
-def test_saver_basic():
-    covid19 = epiworld.ModelSEIRCONN(
-        name              = 'covid-19',
-        n                 = 10000,
-        prevalence        = .01,
-        contact_rate      = 2.0,
-        transmission_rate = .1,
-        incubation_days   = 7.0,
-        recovery_rate     = 0.14
-    )
-
-    saver = epiworld.Saver("total_hist", "virus_hist")
-
-    saver.run_multiple(covid19, 100, 4, nthreads=1)
-
-    # TODO: Verify things worked correctly, as is the point of tesing.