clarify BaseAgentMode operation

teeth_agent/agent.py

@@ -158,12 +158,9 @@ class TeethAgent(object):
                 if not last_command.is_done():
                     raise errors.CommandExecutionError('agent is busy')
 
-            if command_name not in self.mode_implementation:
-                raise errors.InvalidCommandError(command_name)
-
             try:
-                command_fn = self.mode_implementation[command_name]
-                result = command_fn(command_name, **kwargs)
+                result = self.mode_implementation.execute(command_name,
+                                                          **kwargs)
                 if not isinstance(result, base.BaseCommandResult):
                     result = base.SyncCommandResult(command_name,
                                                     kwargs,

teeth_agent/base.py

@@ -118,8 +118,15 @@ class AsyncCommandResult(BaseCommandResult):
         pass
 
 
-class BaseAgentMode(dict):
+class BaseAgentMode(object):
     def __init__(self, name):
         super(BaseAgentMode, self).__init__()
         self.log = structlog.get_logger(agent_mode=name)
         self.name = name
+        self.command_map = {}
+
+    def execute(self, command_name, **kwargs):
+        if command_name not in self.command_map:
+            raise errors.InvalidCommandError(command_name)
+
+        return self.command_map[command_name](command_name, **kwargs)

teeth_agent/standby.py

@@ -135,9 +135,9 @@ class RunImageCommand(base.AsyncCommandResult):
 class StandbyMode(base.BaseAgentMode):
     def __init__(self):
         super(StandbyMode, self).__init__('STANDBY')
-        self['cache_images'] = self.cache_images
-        self['prepare_image'] = self.prepare_image
-        self['run_image'] = self.run_image
+        self.command_map['cache_images'] = self.cache_images
+        self.command_map['prepare_image'] = self.prepare_image
+        self.command_map['run_image'] = self.run_image
 
     def _validate_image_info(self, image_info):
         for field in ['id', 'urls', 'hashes']:

teeth_agent/tests/agent.py

@@ -144,7 +144,8 @@ class TestBaseAgent(unittest.TestCase):
 
     def test_execute_command(self):
         do_something_impl = mock.Mock()
-        self.agent.mode_implementation['do_something'] = do_something_impl
+        command_map = self.agent.mode_implementation.command_map
+        command_map['do_something'] = do_something_impl
 
         self.agent.execute_command('do_something', foo='bar')
         do_something_impl.assert_called_once_with('do_something', foo='bar')