From b24bc52223b4d0e97f6c26aef6293539947fef64 Mon Sep 17 00:00:00 2001
From: Ewout ter Hoeven <E.M.terHoeven@student.tudelft.nl>
Date: Sat, 21 Sep 2024 09:23:24 +0200
Subject: [PATCH] Add tests for shuffle_do()

---
 tests/test_agent.py | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/tests/test_agent.py b/tests/test_agent.py
index f43a80c84ec..958b84faaff 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -437,6 +437,40 @@ def test_agentset_map_callable():
     assert all(i == entry for i, entry in zip(results, range(1, 11)))
 
 
+def test_agentset_shuffle_do():
+    """Test AgentSet.shuffle_do method."""
+    model = Model()
+
+    class TestAgentShuffleDo(Agent):
+        def __init__(self, model):
+            super().__init__(model)
+            self.called = False
+
+        def test_method(self):
+            self.called = True
+
+    agents = [TestAgentShuffleDo(model) for _ in range(100)]
+    agentset = AgentSet(agents, model)
+
+    # Test shuffle_do with a string method name
+    agentset.shuffle_do("test_method")
+    assert all(agent.called for agent in agents)
+
+    # Reset the called flag
+    for agent in agents:
+        agent.called = False
+
+    # Test shuffle_do with a callable
+    agentset.shuffle_do(lambda agent: setattr(agent, 'called', True))
+    assert all(agent.called for agent in agents)
+
+    # Verify that the order is indeed shuffled
+    original_order = list(agentset)
+    shuffled_order = []
+    agentset.shuffle_do(lambda agent: shuffled_order.append(agent))
+    assert original_order != shuffled_order, "The order should be different after shuffle_do"
+
+
 def test_agentset_get_attribute():
     """Test AgentSet.get for attributes."""
     model = Model()