From 6cbd507ba111f7a9e5a393caa06ffce4d17afc45 Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Fri, 29 Jul 2016 09:58:42 -0400
Subject: [PATCH 01/13] Make scheduler keep track of dict instead of list

---
 mesa/time.py | 42 +++++++++++++++++++++++-------------------
 1 file changed, 23 insertions(+), 19 deletions(-)

diff --git a/mesa/time.py b/mesa/time.py
index 6083b8a7e8d..8decdce51c3 100644
--- a/mesa/time.py
+++ b/mesa/time.py
@@ -41,14 +41,14 @@ class BaseScheduler:
     model = None
     steps = 0
     time = 0
-    agents = []
+    agents = {}
 
     def __init__(self, model):
         """ Create a new, empty BaseScheduler. """
         self.model = model
         self.steps = 0
         self.time = 0
-        self.agents = []
+        self.agents = {}
 
     def add(self, agent):
         """ Add an Agent object to the schedule.
@@ -58,7 +58,7 @@ def add(self, agent):
             have a step() method.
 
         """
-        self.agents.append(agent)
+        self.agents[agent.unique_id] = agent
 
     def remove(self, agent):
         """ Remove all instances of a given agent from the schedule.
@@ -67,19 +67,19 @@ def remove(self, agent):
             agent: An agent object.
 
         """
-        while agent in self.agents:
-            self.agents.remove(agent)
+        del self.agents[agent.unique_id]
 
     def step(self):
         """ Execute the step of all the agents, one at a time. """
-        for agent in self.agents:
-            agent.step()
+        agent_keys = list(self.agents.keys())
+        for agent_key in agent_keys:
+            self.agents[agent_key].step(self.model)
         self.steps += 1
         self.time += 1
 
     def get_agent_count(self):
         """ Returns the current number of agents in the queue. """
-        return len(self.agents)
+        return len(self.agents.keys())
 
 
 class RandomActivation(BaseScheduler):
@@ -97,9 +97,11 @@ def step(self):
         random order.
 
         """
-        random.shuffle(self.agents)
-        for agent in self.agents:
-            agent.step()
+        agent_keys = list(self.agents.keys())
+        random.shuffle(agent_keys)
+        
+        for agent_key in agent_keys:
+            self.agents[agent_key].step(self.model)
         self.steps += 1
         self.time += 1
 
@@ -114,10 +116,11 @@ class SimultaneousActivation(BaseScheduler):
     """
     def step(self):
         """ Step all agents, then advance them. """
-        for agent in self.agents:
-            agent.step()
-        for agent in self.agents:
-            agent.advance()
+        agent_keys = list(self.agents.keys())
+        for agent_key in agent_keys:
+            self.agents[agent_key].step(self.model)
+        for agent_key in agent_keys:
+            self.agents[agent_key].advance(self.model)
         self.steps += 1
         self.time += 1
 
@@ -161,13 +164,14 @@ def __init__(self, model, stage_list=["step"], shuffle=False,
 
     def step(self):
         """ Executes all the stages for all agents. """
+        agent_keys = list(self.agents.keys())
         if self.shuffle:
-            random.shuffle(self.agents)
+            random.shuffle(agent_keys)
         for stage in self.stage_list:
-            for agent in self.agents:
-                getattr(agent, stage)()  # Run stage
+            for agent_key in agent_keys:
+                getattr(self.agents[agent_key], stage)()  # Run stage
             if self.shuffle_between_stages:
-                random.shuffle(self.agents)
+                random.shuffle(agent_keys)
             self.time += self.stage_time
 
         self.steps += 1

From cd5555fd9014bad16fe651169b0506a9f4201910 Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Fri, 29 Jul 2016 09:59:07 -0400
Subject: [PATCH 02/13] Add method to keep track of unique IDs

---
 mesa/model.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/mesa/model.py b/mesa/model.py
index 6e4d4241157..760257aa5df 100644
--- a/mesa/model.py
+++ b/mesa/model.py
@@ -30,6 +30,7 @@ def __init__(self, seed=None):
         random.seed(seed)
         self.running = True
         self.schedule = None
+        self.current_id = 0
 
     def run_model(self):
         """ Run the model until the end condition is reached. Overload as
@@ -42,3 +43,8 @@ def run_model(self):
     def step(self):
         """ A single step. Fill in here. """
         pass
+
+    def next_id(self):
+        """ Return the next unique ID for agents, increment current_id"""
+        self.current_id += 1
+        return self.current_id

From 44f92dcc5cd145c083e039e848808b034ac2b877 Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Fri, 29 Jul 2016 10:00:01 -0400
Subject: [PATCH 03/13] Update datacollector to handle agent dict

---
 mesa/datacollection.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mesa/datacollection.py b/mesa/datacollection.py
index e923d129050..3c96a51a1b2 100644
--- a/mesa/datacollection.py
+++ b/mesa/datacollection.py
@@ -139,7 +139,7 @@ def collect(self, model):
         if self.agent_reporters:
             for var, reporter in self.agent_reporters.items():
                 agent_records = []
-                for agent in model.schedule.agents:
+                for agent in model.schedule.agents.values():
                     agent_records.append((agent.unique_id, reporter(agent)))
                 self.agent_vars[var].append(agent_records)
 

From 710e673370e1fdaa46e73e4c1b1b0019eda6ec53 Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Fri, 29 Jul 2016 10:00:19 -0400
Subject: [PATCH 04/13] Update WolfSheep model to handle agent dict

---
 examples/WolfSheep/wolf_sheep/agents.py      | 19 +++++++++-------
 examples/WolfSheep/wolf_sheep/model.py       |  9 ++++----
 examples/WolfSheep/wolf_sheep/random_walk.py |  4 ++--
 examples/WolfSheep/wolf_sheep/schedule.py    | 24 +++++++++-----------
 4 files changed, 29 insertions(+), 27 deletions(-)

diff --git a/examples/WolfSheep/wolf_sheep/agents.py b/examples/WolfSheep/wolf_sheep/agents.py
index 7c4330a9657..b2c1c27fb3e 100644
--- a/examples/WolfSheep/wolf_sheep/agents.py
+++ b/examples/WolfSheep/wolf_sheep/agents.py
@@ -14,8 +14,8 @@ class Sheep(RandomWalker):
 
     energy = None
 
-    def __init__(self, pos, model, moore, energy=None):
-        super().__init__(pos, model, moore=moore)
+    def __init__(self, unique_id, pos, model, moore, energy=None):
+        super().__init__(unique_id, pos, model, moore=moore)
         self.energy = energy
 
     def step(self):
@@ -47,7 +47,8 @@ def step(self):
             # Create a new sheep:
             if self.model.grass:
                 self.energy /= 2
-            lamb = Sheep(self.pos, self.model, self.moore, self.energy)
+            lamb = Sheep(self.model.next_id(), self.pos, self.model,
+                         self.moore, self.energy)
             self.model.grid.place_agent(lamb, self.pos)
             self.model.schedule.add(lamb)
 
@@ -59,8 +60,8 @@ class Wolf(RandomWalker):
 
     energy = None
 
-    def __init__(self, pos, model, moore, energy=None):
-        super().__init__(pos, model, moore=moore)
+    def __init__(self, unique_id, pos, model, moore, energy=None):
+        super().__init__(unique_id, pos, model, moore=moore)
         self.energy = energy
 
     def step(self):
@@ -87,7 +88,8 @@ def step(self):
             if random.random() < self.model.wolf_reproduce:
                 # Create a new wolf cub
                 self.energy /= 2
-                cub = Wolf(self.pos, self.model, self.moore, self.energy)
+                cub = Wolf(self.model.next_id(), self.pos, self.model,
+                           self.moore, self.energy)
                 self.model.grid.place_agent(cub, cub.pos)
                 self.model.schedule.add(cub)
 
@@ -97,7 +99,7 @@ class GrassPatch(Agent):
     A patch of grass that grows at a fixed rate and it is eaten by sheep
     '''
 
-    def __init__(self, pos, model, fully_grown, countdown):
+    def __init__(self, unique_id, pos, model, fully_grown, countdown):
         '''
         Creates a new patch of grass
 
@@ -105,9 +107,10 @@ def __init__(self, pos, model, fully_grown, countdown):
             grown: (boolean) Whether the patch of grass is fully grown or not
             countdown: Time for the patch of grass to be fully grown again
         '''
-        super().__init__(pos, model)
+        super().__init__(unique_id, model)
         self.fully_grown = fully_grown
         self.countdown = countdown
+        self.pos = pos
 
     def step(self):
         if not self.fully_grown:
diff --git a/examples/WolfSheep/wolf_sheep/model.py b/examples/WolfSheep/wolf_sheep/model.py
index 7b79cc80263..5b93dc57167 100644
--- a/examples/WolfSheep/wolf_sheep/model.py
+++ b/examples/WolfSheep/wolf_sheep/model.py
@@ -60,7 +60,7 @@ def __init__(self, height=20, width=20,
                                  once it is eaten
             sheep_gain_from_food: Energy sheep gain from grass, if enabled.
         '''
-
+        super().__init__()
         # Set parameters
         self.height = height
         self.width = width
@@ -84,7 +84,7 @@ def __init__(self, height=20, width=20,
             x = random.randrange(self.width)
             y = random.randrange(self.height)
             energy = random.randrange(2 * self.sheep_gain_from_food)
-            sheep = Sheep((x, y), self, True, energy)
+            sheep = Sheep(self.next_id(), (x, y), self, True, energy)
             self.grid.place_agent(sheep, (x, y))
             self.schedule.add(sheep)
 
@@ -93,7 +93,7 @@ def __init__(self, height=20, width=20,
             x = random.randrange(self.width)
             y = random.randrange(self.height)
             energy = random.randrange(2 * self.wolf_gain_from_food)
-            wolf = Wolf((x, y), self, True, energy)
+            wolf = Wolf(self.next_id(), (x, y), self, True, energy)
             self.grid.place_agent(wolf, (x, y))
             self.schedule.add(wolf)
 
@@ -108,7 +108,8 @@ def __init__(self, height=20, width=20,
                 else:
                     countdown = random.randrange(self.grass_regrowth_time)
 
-                patch = GrassPatch((x, y), self, fully_grown, countdown)
+                patch = GrassPatch(self.next_id(), (x, y), self,
+                                   fully_grown, countdown)
                 self.grid.place_agent(patch, (x, y))
                 self.schedule.add(patch)
 
diff --git a/examples/WolfSheep/wolf_sheep/random_walk.py b/examples/WolfSheep/wolf_sheep/random_walk.py
index d78b47a55f2..c034cc9bde5 100644
--- a/examples/WolfSheep/wolf_sheep/random_walk.py
+++ b/examples/WolfSheep/wolf_sheep/random_walk.py
@@ -21,7 +21,7 @@ class RandomWalker(Agent):
     y = None
     moore = True
 
-    def __init__(self, pos, model, moore=True):
+    def __init__(self, unique_id, pos, model, moore=True):
         '''
         grid: The MultiGrid object in which the agent lives.
         x: The agent's current x coordinate
@@ -29,7 +29,7 @@ def __init__(self, pos, model, moore=True):
         moore: If True, may move in all 8 directions.
                 Otherwise, only up, down, left, right.
         '''
-        super().__init__(pos, model)
+        super().__init__(unique_id, model)
         self.pos = pos
         self.moore = moore
 
diff --git a/examples/WolfSheep/wolf_sheep/schedule.py b/examples/WolfSheep/wolf_sheep/schedule.py
index 9b59a5e2483..eb6efa1a60c 100644
--- a/examples/WolfSheep/wolf_sheep/schedule.py
+++ b/examples/WolfSheep/wolf_sheep/schedule.py
@@ -14,11 +14,11 @@ class RandomActivationByBreed(RandomActivation):
 
     Assumes that all agents have a step() method.
     '''
-    agents_by_breed = defaultdict(list)
+    agents_by_breed = defaultdict(dict)
 
     def __init__(self, model):
         super().__init__(model)
-        self.agents_by_breed = defaultdict(list)
+        self.agents_by_breed = defaultdict(dict)
 
     def add(self, agent):
         '''
@@ -28,21 +28,19 @@ def add(self, agent):
             agent: An Agent to be added to the schedule.
         '''
 
-        self.agents.append(agent)
+        self.agents[agent.unique_id] = agent
         agent_class = type(agent)
-        self.agents_by_breed[agent_class].append(agent)
+        self.agents_by_breed[agent_class][agent.unique_id] = agent
 
     def remove(self, agent):
         '''
         Remove all instances of a given agent from the schedule.
         '''
 
-        while agent in self.agents:
-            self.agents.remove(agent)
+        del self.agents[agent.unique_id]
 
         agent_class = type(agent)
-        while agent in self.agents_by_breed[agent_class]:
-            self.agents_by_breed[agent_class].remove(agent)
+        del self.agents_by_breed[agent_class][agent.unique_id]
 
     def step(self, by_breed=True):
         '''
@@ -67,13 +65,13 @@ def step_breed(self, breed):
         Args:
             breed: Class object of the breed to run.
         '''
-        agents = self.agents_by_breed[breed]
-        random.shuffle(agents)
-        for agent in agents:
-            agent.step()
+        agent_keys = list(self.agents_by_breed[breed].keys())
+        random.shuffle(agent_keys)
+        for agent_key in agent_keys:
+            self.agents_by_breed[breed][agent_key].step()
 
     def get_breed_count(self, breed_class):
         '''
         Returns the current number of agents of certain breed in the queue.
         '''
-        return len(self.agents_by_breed[breed_class])
+        return len(self.agents_by_breed[breed_class].values())

From bccd3df1efd419462d2a51e0f03532ee9556be10 Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Fri, 29 Jul 2016 13:07:51 -0400
Subject: [PATCH 05/13] Added test for lifespan of agents

---
 tests/test_lifespan.py | 96 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 96 insertions(+)
 create mode 100644 tests/test_lifespan.py

diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py
new file mode 100644
index 00000000000..0f8c7efe303
--- /dev/null
+++ b/tests/test_lifespan.py
@@ -0,0 +1,96 @@
+import unittest
+
+from mesa.time import RandomActivation
+from mesa.datacollection import DataCollector
+from mesa import *
+import numpy as np
+
+class LifeTimeModel(Model):
+    '''Simple model for running models with a finite life'''
+    def __init__(self, agent_lifetime = 1, n_agents = 10):
+        super().__init__()
+        
+        self.agent_lifetime = agent_lifetime
+        self.n_agents = n_agents
+
+        ## keep track of the the remaining life of an agent and
+        ## how many ticks it has seen
+        self.datacollector = DataCollector(
+            agent_reporters = {"remaining_life" : lambda a: a.remaining_life,
+                               "steps" : lambda a: a.steps})
+        
+        self.current_ID = 0
+        self.schedule = RandomActivation(self)
+        
+        for _ in range(n_agents):
+            self.schedule.add(FiniteLifeAgent(self.next_id(),
+                                              self.agent_lifetime,
+                                              self))
+
+    def step(self):
+        '''Add agents back to n_agents in each step'''
+        
+        self.datacollector.collect(self)
+        self.schedule.step()
+        
+        if len(self.schedule.agents) < self.n_agents:
+            for _ in range(self.n_agents - len(self.schedule.agents)):
+                self.schedule.add(FiniteLifeAgent(self.next_id(),
+                                                  self.agent_lifetime,
+                                                  self))
+
+    def run_model(self, step_count = 100):
+        for _ in range(step_count):
+            self.step()
+
+class FiniteLifeAgent(Agent):
+    
+    '''An agent that is supposed to live for a finite number of ticks.
+    Also has a 10% chance of dying in each tick.
+
+    '''
+    
+    def __init__(self, unique_id, lifetime, model):
+        super().__init__(unique_id, model)
+        self.remaining_life = lifetime
+        self.steps = 0
+
+    def step(self, model):
+        inactivated = self.inactivate(model)
+        if not inactivated:
+            self.steps += 1 # keep track of how many ticks are seen
+            if np.random.binomial(1,0.1) is not 0: # 10% chance of dying
+                model.schedule.remove(self)
+            
+    def inactivate(self, model):
+        self.remaining_life -= 1
+        if self.remaining_life < 0:
+            model.schedule.remove(self)
+            return True
+        return False
+
+
+class TestAgentLifespan(unittest.TestCase):
+    def setUp(self):
+        self.model = LifeTimeModel()
+        self.model.run_model()
+        self.df = self.model.datacollector.get_agent_vars_dataframe()
+        self.df = self.df.reset_index()
+        
+    def test_ticks_seen(self):
+        '''Each agent should be activated no more than one time'''
+        assert self.df.steps.max() == 1
+
+    def test_agent_lifetime(self):
+        lifetimes = self.df.groupby(["AgentID"]).Step.agg({"Step" : lambda x: len(x)})
+        print( lifetimes.Step.max() )
+        assert lifetimes.Step.max() == 2
+
+
+if __name__ == '__main__':
+    unittest.main()
+
+        
+        
+
+    

From 19f34b46c268b326012ae6b669f090ee4d483762 Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Fri, 29 Jul 2016 13:08:16 -0400
Subject: [PATCH 06/13] Fix bug related to agent dict

---
 mesa/time.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/mesa/time.py b/mesa/time.py
index 8decdce51c3..195a7bf52c8 100644
--- a/mesa/time.py
+++ b/mesa/time.py
@@ -73,7 +73,7 @@ def step(self):
         """ Execute the step of all the agents, one at a time. """
         agent_keys = list(self.agents.keys())
         for agent_key in agent_keys:
-            self.agents[agent_key].step(self.model)
+            self.agents[agent_key].step()
         self.steps += 1
         self.time += 1
 
@@ -101,7 +101,7 @@ def step(self):
         random.shuffle(agent_keys)
         
         for agent_key in agent_keys:
-            self.agents[agent_key].step(self.model)
+            self.agents[agent_key].step()
         self.steps += 1
         self.time += 1
 
@@ -118,9 +118,9 @@ def step(self):
         """ Step all agents, then advance them. """
         agent_keys = list(self.agents.keys())
         for agent_key in agent_keys:
-            self.agents[agent_key].step(self.model)
+            self.agents[agent_key].step()
         for agent_key in agent_keys:
-            self.agents[agent_key].advance(self.model)
+            self.agents[agent_key].advance()
         self.steps += 1
         self.time += 1
 

From 386f90d9ab4391688af814fac8f47c7f843ea5b3 Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Fri, 29 Jul 2016 13:08:44 -0400
Subject: [PATCH 07/13] Make batchrunner compatible with agent dict

---
 mesa/batchrunner.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mesa/batchrunner.py b/mesa/batchrunner.py
index 999d72487b2..efd302312bf 100644
--- a/mesa/batchrunner.py
+++ b/mesa/batchrunner.py
@@ -103,7 +103,7 @@ def collect_model_vars(self, model):
     def collect_agent_vars(self, model):
         """ Run reporters and collect agent-level variables. """
         agent_vars = {}
-        for agent in model.schedule.agents:
+        for agent in model.schedule.agents.values():
             agent_record = {}
             for var, reporter in self.agent_reporters.items():
                 agent_record[var] = reporter(agent)

From 787f818277d561ad1730ca127c9744f857ba8e8b Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Fri, 29 Jul 2016 13:09:22 -0400
Subject: [PATCH 08/13] Added call for unittest. Fix datacollector test

---
 tests/test_batchrunner.py   | 3 +++
 tests/test_datacollector.py | 5 ++++-
 tests/test_grid.py          | 3 +++
 tests/test_space.py         | 3 +++
 4 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/tests/test_batchrunner.py b/tests/test_batchrunner.py
index 2e9242c4c7b..0cf9b238966 100644
--- a/tests/test_batchrunner.py
+++ b/tests/test_batchrunner.py
@@ -92,3 +92,6 @@ def test_agent_level_vars(self):
             len(self.params['model_param']) * \
             self.iterations
         assert agent_vars.shape == (rows, 6)
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_datacollector.py b/tests/test_datacollector.py
index 3c62f19cd88..b08da7ff8c8 100644
--- a/tests/test_datacollector.py
+++ b/tests/test_datacollector.py
@@ -61,7 +61,7 @@ def setUp(self):
         for i in range(7):
             self.model.step()
         # Write to table:
-        for agent in self.model.schedule.agents:
+        for agent in self.model.schedule.agents.values():
             agent.write_final_values()
 
     def test_model_vars(self):
@@ -117,3 +117,6 @@ def test_exports(self):
 
         with self.assertRaises(Exception):
             table_df = data_collector.get_table_dataframe("not a real table")
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_grid.py b/tests/test_grid.py
index e2fd7f27fcf..abd6c379913 100644
--- a/tests/test_grid.py
+++ b/tests/test_grid.py
@@ -317,3 +317,6 @@ def test_neighbors(self):
 
         neighbors = self.grid.get_neighbors((1, 3), moore=False, radius=2)
         assert len(neighbors) == 11
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_space.py b/tests/test_space.py
index d490e50857b..5cbf56c6103 100644
--- a/tests/test_space.py
+++ b/tests/test_space.py
@@ -100,3 +100,6 @@ def test_neighborhood_retrieval(self):
 
         neighbors_3 = self.space.get_neighbors((-30, -30), 10)
         assert len(neighbors_3) == 0
+
+if __name__ == '__main__':
+    unittest.main()

From dbf492fb1b447613d2f630486e08d1e5adc8c3f6 Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Fri, 29 Jul 2016 13:16:46 -0400
Subject: [PATCH 09/13] Re-wrote agent activation unittests

---
 tests/test_time.py | 44 ++++++++++++++++++++++++++++----------------
 1 file changed, 28 insertions(+), 16 deletions(-)

diff --git a/tests/test_time.py b/tests/test_time.py
index b97caebae4e..a3e04cbd92a 100644
--- a/tests/test_time.py
+++ b/tests/test_time.py
@@ -2,6 +2,7 @@
 Test the advanced schedulers.
 '''
 
+import unittest
 from unittest import TestCase
 from unittest.mock import patch
 from mesa import Model, Agent
@@ -17,7 +18,11 @@ class MockAgent(Agent):
     '''
     Minimalistic agent for testing purposes.
     '''
-
+    def __init__(self, unique_id, model):
+        super().__init__(unique_id, model)
+        self.steps = 0
+        self.advances = 0
+    
     def stage_one(self):
         self.model.log.append(self.unique_id + "_1")
 
@@ -25,8 +30,10 @@ def stage_two(self):
         self.model.log.append(self.unique_id + "_2")
 
     def advance(self):
-        pass
+        self.advances += 1
 
+    def step(self):
+        self.steps += 1
 
 class MockModel(Model):
 
@@ -106,8 +113,9 @@ def test_remove(self):
         Test staged activation can remove an agent
         '''
         model = MockModel(shuffle=True)
-        agent = model.schedule.agents[0]
-        model.schedule.remove(model.schedule.agents[0])
+        agent_keys = list(model.schedule.agents.keys())
+        agent = model.schedule.agents[agent_keys[0]]
+        model.schedule.remove(agent)
         assert agent not in model.schedule.agents
 
 
@@ -141,11 +149,11 @@ def test_random_activation_step_steps_each_agent(self):
         Test the random activation step causes each agent to step
         '''
 
-        with patch('test_time.MockAgent.step') as mock_agent_step:
-            model = MockModel(activation=RANDOM)
-            model.step()
-            # one step for each of 2 agents
-            assert mock_agent_step.call_count == 2
+        model = MockModel(activation=RANDOM)
+        model.step()
+        agent_steps = [i.steps for i in model.schedule.agents.values()]
+        # one step for each of 2 agents
+        assert all(map(lambda x: x == 1, agent_steps))
 
 
 class TestSimultaneousActivation(TestCase):
@@ -158,10 +166,14 @@ def test_simultaneous_activation_step_steps_and_advances_each_agent(self):
         Test the simultaneous activation step causes each agent to step
         '''
 
-        with patch('test_time.MockAgent.step') as mock_agent_step,\
-                patch('test_time.MockAgent.advance') as mock_agent_advance:
-            model = MockModel(activation=SIMULTANEOUS)
-            model.step()
-            # one step for each of 2 agents
-            assert mock_agent_step.call_count == 2
-            assert mock_agent_advance.call_count == 2
+        
+        model = MockModel(activation=SIMULTANEOUS)
+        model.step()
+        # one step for each of 2 agents
+        agent_steps = [i.steps for i in model.schedule.agents.values()]
+        agent_advances = [i.advances for i in model.schedule.agents.values()]
+        assert all(map(lambda x: x == 1, agent_steps))
+        assert all(map(lambda x: x == 1, agent_advances))
+
+if __name__ == '__main__':
+    unittest.main()

From 550665bac38c1ad7da53f30cbd8c40cfa4e76dda Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Fri, 29 Jul 2016 14:01:16 -0400
Subject: [PATCH 10/13] Fix lifespan test

---
 tests/test_lifespan.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py
index 0f8c7efe303..ce786f90f3b 100644
--- a/tests/test_lifespan.py
+++ b/tests/test_lifespan.py
@@ -54,18 +54,19 @@ def __init__(self, unique_id, lifetime, model):
         super().__init__(unique_id, model)
         self.remaining_life = lifetime
         self.steps = 0
+        self.model = model
 
-    def step(self, model):
-        inactivated = self.inactivate(model)
+    def step(self):
+        inactivated = self.inactivate()
         if not inactivated:
             self.steps += 1 # keep track of how many ticks are seen
             if np.random.binomial(1,0.1) is not 0: # 10% chance of dying
-                model.schedule.remove(self)
+                self.model.schedule.remove(self)
             
-    def inactivate(self, model):
+    def inactivate(self):
         self.remaining_life -= 1
         if self.remaining_life < 0:
-            model.schedule.remove(self)
+            self.model.schedule.remove(self)
             return True
         return False
 

From 3dd085ccf6993b24b45921bd034007a7d841e237 Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Fri, 29 Jul 2016 14:13:36 -0400
Subject: [PATCH 11/13] Fix time test to work with dict keys

---
 tests/test_time.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/test_time.py b/tests/test_time.py
index a3e04cbd92a..86e29e5ce14 100644
--- a/tests/test_time.py
+++ b/tests/test_time.py
@@ -88,7 +88,8 @@ def test_no_shuffle(self):
         '''
         model = MockModel(shuffle=False)
         model.step()
-        assert model.log == self.expected_output
+        model.step()
+        assert all([i==j for i,j in zip(model.log[:4], model.log[4:])])
 
     def test_shuffle(self):
         '''

From 311273003e288a727c755cf25f61e387981dbed5 Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Fri, 29 Jul 2016 19:25:05 -0400
Subject: [PATCH 12/13] Fix PEP8

---
 tests/test_lifespan.py | 46 +++++++++++++++++-------------------------
 tests/test_time.py     |  7 +++----
 2 files changed, 22 insertions(+), 31 deletions(-)

diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py
index ce786f90f3b..20c432cf534 100644
--- a/tests/test_lifespan.py
+++ b/tests/test_lifespan.py
@@ -2,26 +2,27 @@
 
 from mesa.time import RandomActivation
 from mesa.datacollection import DataCollector
-from mesa import *
+from mesa import Model, Agent
 import numpy as np
 
+
 class LifeTimeModel(Model):
     '''Simple model for running models with a finite life'''
-    def __init__(self, agent_lifetime = 1, n_agents = 10):
+    def __init__(self, agent_lifetime=1, n_agents=10):
         super().__init__()
-        
+
         self.agent_lifetime = agent_lifetime
         self.n_agents = n_agents
 
-        ## keep track of the the remaining life of an agent and
-        ## how many ticks it has seen
+        # keep track of the the remaining life of an agent and
+        # how many ticks it has seen
         self.datacollector = DataCollector(
-            agent_reporters = {"remaining_life" : lambda a: a.remaining_life,
-                               "steps" : lambda a: a.steps})
-        
+            agent_reporters={"remaining_life": lambda a: a.remaining_life,
+                             "steps": lambda a: a.steps})
+
         self.current_ID = 0
         self.schedule = RandomActivation(self)
-        
+
         for _ in range(n_agents):
             self.schedule.add(FiniteLifeAgent(self.next_id(),
                                               self.agent_lifetime,
@@ -29,27 +30,24 @@ def __init__(self, agent_lifetime = 1, n_agents = 10):
 
     def step(self):
         '''Add agents back to n_agents in each step'''
-        
         self.datacollector.collect(self)
         self.schedule.step()
-        
+
         if len(self.schedule.agents) < self.n_agents:
             for _ in range(self.n_agents - len(self.schedule.agents)):
                 self.schedule.add(FiniteLifeAgent(self.next_id(),
                                                   self.agent_lifetime,
                                                   self))
 
-    def run_model(self, step_count = 100):
+    def run_model(self, step_count=100):
         for _ in range(step_count):
             self.step()
 
+
 class FiniteLifeAgent(Agent):
-    
     '''An agent that is supposed to live for a finite number of ticks.
     Also has a 10% chance of dying in each tick.
-
     '''
-    
     def __init__(self, unique_id, lifetime, model):
         super().__init__(unique_id, model)
         self.remaining_life = lifetime
@@ -59,10 +57,10 @@ def __init__(self, unique_id, lifetime, model):
     def step(self):
         inactivated = self.inactivate()
         if not inactivated:
-            self.steps += 1 # keep track of how many ticks are seen
-            if np.random.binomial(1,0.1) is not 0: # 10% chance of dying
+            self.steps += 1  # keep track of how many ticks are seen
+            if np.random.binomial(1, 0.1) is not 0:  # 10% chance of dying
                 self.model.schedule.remove(self)
-            
+
     def inactivate(self):
         self.remaining_life -= 1
         if self.remaining_life < 0:
@@ -77,21 +75,15 @@ def setUp(self):
         self.model.run_model()
         self.df = self.model.datacollector.get_agent_vars_dataframe()
         self.df = self.df.reset_index()
-        
+
     def test_ticks_seen(self):
         '''Each agent should be activated no more than one time'''
         assert self.df.steps.max() == 1
 
     def test_agent_lifetime(self):
-        lifetimes = self.df.groupby(["AgentID"]).Step.agg({"Step" : lambda x: len(x)})
-        print( lifetimes.Step.max() )
+        lifetimes = self.df.groupby(["AgentID"]).Step.agg({"Step":
+                                                           lambda x: len(x)})
         assert lifetimes.Step.max() == 2
 
-
 if __name__ == '__main__':
     unittest.main()
-
-        
-        
-
-    
diff --git a/tests/test_time.py b/tests/test_time.py
index 86e29e5ce14..e75b21579e3 100644
--- a/tests/test_time.py
+++ b/tests/test_time.py
@@ -22,7 +22,7 @@ def __init__(self, unique_id, model):
         super().__init__(unique_id, model)
         self.steps = 0
         self.advances = 0
-    
+
     def stage_one(self):
         self.model.log.append(self.unique_id + "_1")
 
@@ -35,6 +35,7 @@ def advance(self):
     def step(self):
         self.steps += 1
 
+
 class MockModel(Model):
 
     def __init__(self, shuffle=False, activation=STAGED):
@@ -89,7 +90,7 @@ def test_no_shuffle(self):
         model = MockModel(shuffle=False)
         model.step()
         model.step()
-        assert all([i==j for i,j in zip(model.log[:4], model.log[4:])])
+        assert all([i == j for i, j in zip(model.log[:4], model.log[4:])])
 
     def test_shuffle(self):
         '''
@@ -166,8 +167,6 @@ def test_simultaneous_activation_step_steps_and_advances_each_agent(self):
         '''
         Test the simultaneous activation step causes each agent to step
         '''
-
-        
         model = MockModel(activation=SIMULTANEOUS)
         model.step()
         # one step for each of 2 agents

From e066756aead1fbc3ebd90c22fa8f87f09a21a182 Mon Sep 17 00:00:00 2001
From: Colin Maxwell <cs.maxwell@gmail.com>
Date: Mon, 1 Aug 2016 12:58:30 -0400
Subject: [PATCH 13/13] Fix more PEP8

---
 mesa/time.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mesa/time.py b/mesa/time.py
index 195a7bf52c8..d4bd5a11386 100644
--- a/mesa/time.py
+++ b/mesa/time.py
@@ -99,7 +99,7 @@ def step(self):
         """
         agent_keys = list(self.agents.keys())
         random.shuffle(agent_keys)
-        
+
         for agent_key in agent_keys:
             self.agents[agent_key].step()
         self.steps += 1