--- a/mamba/habitats/editor.py	Fri Sep 23 22:17:46 2011 +0200
+++ b/mamba/habitats/editor.py	Thu Sep 22 17:32:38 2011 +0200
@@ -74,6 +74,7 @@
         self.sprite_mode = 'Add'
         self.sprite_cls = None
         # Map for test mode lookups
+        self.replays = {}
         self.action_map = {
                 K_UP: UP,
                 K_DOWN: DOWN,
@@ -213,6 +214,22 @@
                     tile_button.add_callback('clicked', self.sprite_tool,
                             name, None, None)
                     tool_list.append(tile_button)
+            elif self.mode == 'Test':
+                for replay in [1, 2]:
+                    store_button = TextButton((0, 0),
+                            'Store run as %s' % replay)
+                    store_button.add_callback('clicked', self.store_replay,
+                            replay)
+                    tool_list.append(store_button)
+                    replay_button = TextButton((0, 0),
+                            'Replay run %s' % replay)
+                    replay_button.add_callback('clicked', self.do_replay,
+                            replay)
+                    tool_list.append(replay_button)
+                last_button = TextButton((0, 0),
+                        'Replay last run')
+                last_button.add_callback('clicked', self.do_replay, None)
+                tool_list.append(last_button)
             self.tool_widget = ToolListWidget((button_left, button_height),
                     tool_list, MAX_TOOLS, start_key=K_2)
             self.toolbar[self.mode] = self.tool_widget
@@ -396,6 +413,7 @@
         self.edit_widget = EditLevelWidget(self.level)
         self.edit_widget.add_callback(SnakeDiedEvent, self.halt_test)
         self.edit_widget.add_callback(LevelCompletedEvent, self.halt_test)
+        self.replays = {}
         self.container.add(self.edit_widget)
         self.clear_toolbar()
         self.setup_toolbar()
@@ -546,3 +564,17 @@
             self.level.replace_sprite(sprite)
         self.level.restart()
         return True
+
+    def store_replay(self, ev, widget, number):
+        self.replays[number] = self.edit_widget.get_replay()
+
+    def do_replay(self, ev, widget, number):
+        if number:
+            # We do nothing if the user tries to replay something thay
+            # isn't there
+            if number in self.replays:
+                # Replay stored replay
+                self.edit_widget.replay(self.replays[number])
+        else:
+            # Replaying last run
+            self.edit_widget.replay()