Add tests for agent removal and model agent_types property

Author: EwoutH

tests/test_agent.py

@@ -0,0 +1,16 @@
+from mesa.agent import Agent
+from mesa.model import Model
+
+
+def test_agent_removal():
+    class TestAgent(Agent):
+        pass
+
+    model = Model()
+    agent = TestAgent(model.next_id(), model)
+    # Check if the agent is added
+    assert agent in model.agents[type(agent)]
+
+    agent.remove()
+    # Check if the agent is removed
+    assert agent not in model.agents[type(agent)]

tests/test_model.py

@@ -1,3 +1,4 @@
+from mesa.agent import Agent
 from mesa.model import Model
 
 
@@ -40,3 +41,13 @@ def test_reset_randomizer(newseed=42):
     assert model._seed == oldseed
     model.reset_randomizer(seed=newseed)
     assert model._seed == newseed
+
+
+def test_agent_types():
+    class TestAgent(Agent):
+        pass
+
+    model = Model()
+    test_agent = TestAgent(model.next_id(), model)
+    assert test_agent in model.agents[type(test_agent)]
+    assert type(test_agent) in model.agent_types