Support agent selection via the run() function

Author: btalb

benchbot_api/benchbot.py

@@ -76,16 +76,11 @@ def __init__(self,
                  supervisor_address='http://' + DEFAULT_ADDRESS + ':' +
                  str(DEFAULT_PORT) + '/',
                  auto_start=True):
-        if agent is not None and not isinstance(agent, Agent):
-            raise ValueError("BenchBot received an agent of type '%s' "
-                             "which is not an instance of '%s'." %
-                             (agent.__class__.__name__, Agent.__name__))
-        self.agent = agent
-
+        self.agent = None
         self.supervisor_address = supervisor_address
         self._connection_callbacks = {}
-        if auto_start:
-            self.start()
+
+        self.set_agent(agent, auto_start)
 
     def _build_address(self, route_name, route_type=RouteType.CONNECTION):
         """Builds an address for communication with a running instance of 
@@ -304,11 +299,13 @@ def results_functions(self):
             for r in self._query('/', BenchBot.RouteType.RESULTS)
         }
 
-    def run(self):
+    def run(self, agent=None):
         """Helper function that runs the robot according to the agent given.
         Generally, you should use this function and implement your object in
         your own custom agent class. 
         """
+        if agent is not None:
+            self.set_agent(agent)
         if self.agent is None:
             raise RuntimeError(
                 "Can't call Benchbot.run() without an agent attached. Either "
@@ -332,6 +329,21 @@ def scene_fn():
         self.agent.save_result(self.result_filename, self.empty_results(),
                                self.results_functions())
 
+    def set_agent(self, agent, auto_start=True):
+        """Updates the current agent, and starts its connection with a BenchBot
+        Supervisor if requested"""
+        if agent is None:
+            self.agent = None
+            return
+
+        if agent is not None and not isinstance(agent, Agent):
+            raise ValueError("BenchBot received an agent of type '%s' "
+                             "which is not an instance of '%s'." %
+                             (agent.__class__.__name__, Agent.__name__))
+        self.agent = agent
+        if auto_start:
+            self.start()
+
     def start(self):
         """Establishes a connect to the Supervisor, and then uses this to
         establish a connection with a running robot. It then initialises all