Update test_time.py

tests/test_time.py

@@ -60,25 +60,29 @@ def __init__(
         super().__init__(seed=seed)
         self.log = []
         self.enable_kill_other_agent = enable_kill_other_agent
-        self.shuffle = shuffle
 
         match activation:
             case "random":
-                pass
+                self.schedule = RandomActivation(self)
             case "staged":
-                pass
+                model_stages = ["stage_one", "model.model_stage", "stage_two"]
+                self.schedule = StagedActivation(
+                    self, stage_list=model_stages, shuffle=shuffle
+                )
             case "simultaneous":
-                pass
+                self.schedule = SimultaneousActivation(self)
             case "random_by_type":
-                pass
+                self.schedule = RandomActivationByType(self)
             case _:
-                raise ValueError(f"Unknown activation type: {activation}")
+                self.schedule = BaseScheduler(self)
 
+        for _ in range(2):
+            agent = MockAgent(self)
+            self.schedule.add(agent)
 
 class MockModel(Model):
-    def __init__(
-        self, seed=None, shuffle=False, activation=STAGED, enable_kill_other_agent=False
-    ):
+
+    def __init__(self, seed=None, shuffle=False, activation=STAGED, enable_kill_other_agent=False):
         """
         Creates a Model instance with a schedule