More bug fixes.

2 files changed, 5 insertions, 4 deletions

diff --git a/oldbot.py b/oldbot.py
index 436375c..e27ed39 100644
--- a/oldbot.py
+++ b/oldbot.py
@@ -23,7 +23,7 @@ class BotWrapper(object):
         proc.stdin.close()
         proc.wait()
 
-        if proc.returncode == 0:
-            output = proc.stdout.read()
+        assert proc.returncode == 0, 'Snake died.'
+        output = proc.stdout.read()
         return output.strip()
 
diff --git a/snake.py b/snake.py
index 46f7db2..e61530c 100755
--- a/snake.py
+++ b/snake.py
@@ -88,7 +88,7 @@ class SnakeEngine(object):
         self.apple = load_image(Sprites.APPLE)
         self.eyes = load_image(Sprites.EYES)
 
-    def add_bot(self, bot):
+    def add_bot(self, bot, colour=None):
         """
         A bot is a callable object, with this method signature:
             def bot_callable(
@@ -103,7 +103,8 @@ class SnakeEngine(object):
         if position is None:
             raise KeyError, "Could not insert snake into the board."
 
-        colour = (randint(0, 255), randint(0, 255), randint(0, 255))
+        if colour is None:
+            colour = (randint(0, 255), randint(0, 255), randint(0, 255))
         self.bots[letter] = [bot, colour, deque([position])]
         return letter