swarming: Add server-side implementation for supplemental bot_config

appengine/swarming/handlers_bot_test.py

 #!/usr/bin/env python
 # coding: utf-8
 # Copyright 2015 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 import base64
 import datetime
 import logging
 import os
@@ skipping 206 matching lines @@
       u'cmd': u'update',
       u'version': old_version,
     }
     self.assertEqual(expected, response)
 
   def test_poll_bot_group_config_change(self):
     params = self.do_handshake()
     params['state']['bot_group_cfg_version'] = 'badversion'
     response = self.post_json('/swarming/api/v1/bot/poll', params)
     expected = {
-      u'cmd': u'restart',
+      u'cmd': u'bot_restart',
       u'message': u'Restarting to pick up new bots.cfg config',
     }
     self.assertEqual(expected, response)
 
   def test_poll_restart(self):
     def mock_should_restart_bot(bot_id, state):
       self.assertEqual('bot1', bot_id)
       expected_state = {
         'bot_group_cfg_version': 'default',
         'running_time': 1234.0,
         'sleep_streak': 0,
         'started_ts': 1410990411.111,
       }
       self.assertEqual(expected_state, state)
       return True, 'Mocked restart message'
     self.mock(bot_management, 'should_restart_bot', mock_should_restart_bot)
 
     params = self.do_handshake()
     response = self.post_json('/swarming/api/v1/bot/poll', params)
     expected = {
-      u'cmd': u'restart',
-      u'message': 'Mocked restart message',
+      u'cmd': u'host_reboot',
+      u'message': u'Mocked restart message',
     }
     self.assertEqual(expected, response)
 
   def test_poll_task_raw(self):
     # Successfully poll a task.
     self.mock(random, 'getrandbits', lambda _: 0x88)
     now = datetime.datetime(2010, 1, 2, 3, 4, 5)
     self.mock_now(now)
     str_now = unicode(now.strftime(DATETIME_FORMAT))
     # A bot polls, gets a task, updates it, completes it.
@@ skipping 168 matching lines @@
   def test_poll_conflicting_dimensions(self):
     params = self.do_handshake()
     self.assertEqual(params['dimensions']['pool'], ['default'])
 
     cfg = bot_groups_config.BotGroupConfig(
         version='default',
         require_luci_machine_token=False,
         require_service_account=None,
         ip_whitelist=None,
         owners=None,
-        dimensions={u'pool': [u'server-side']})
+        dimensions={u'pool': [u'server-side']},
+        bot_config_script=None,
+        bot_config_script_content=None)
     self.mock(bot_auth, 'validate_bot_id_and_fetch_config',
               lambda *args, **kwargs: cfg)
 
     actions = []
     self.mock(stats, 'add_entry', lambda **kwargs: actions.append(kwargs))
 
     # Bot sends 'default' pool, but server config defined it as 'server-side'.
     response = self.post_json('/swarming/api/v1/bot/poll', params)
     self.assertTrue(response.pop(u'duration'))
     expected = {
       u'cmd': u'sleep',
       u'quarantined': False,
     }
     self.assertEqual(expected, response)
 
     # 'server-side' was actually used.
     self.assertEqual([{
         'action': 'bot_active',
         'bot_id': u'bot1',
         'dimensions': {
             u'id': [u'bot1'],
             u'os': [u'Amiga'],
             u'pool': [u'server-side'],
         },
     }], actions)
 
+  def test_poll_extra_bot_config(self):
+    cfg = bot_groups_config.BotGroupConfig(
+        version='default',
+        require_luci_machine_token=False,
+        require_service_account=None,
+        ip_whitelist=None,
+        owners=None,
+        dimensions={},
+        bot_config_script='foo.py',
+        bot_config_script_content='print "Hi";import sys; sys.exit(1)')
+    self.mock(bot_auth, 'validate_bot_id_and_fetch_config',
+              lambda *args, **kwargs: cfg)
+    params = self.do_handshake()
+    self.assertEqual(
+        u'print "Hi";import sys; sys.exit(1)', params['bot_config'])
+
   def test_complete_task_isolated(self):
     # Successfully poll a task.
     self.mock(random, 'getrandbits', lambda _: 0x88)
     now = datetime.datetime(2010, 1, 2, 3, 4, 5)
     self.mock_now(now)
     str_now = unicode(now.strftime(DATETIME_FORMAT))
     # A bot polls, gets a task, updates it, completes it.
     params = self.do_handshake()
     # Enqueue a task.
     _, task_id = self.client_create_task_isolated()
@@ skipping 641 matching lines @@
     self.app.get('/bot_code?tok=%s' % tok, status=200)
 
 
 if __name__ == '__main__':
   if '-v' in sys.argv:
     unittest.TestCase.maxDiff = None
   logging.basicConfig(
       level=logging.DEBUG if '-v' in sys.argv else logging.CRITICAL,
       format='%(levelname)-7s %(filename)s:%(lineno)3d %(message)s')
   unittest.main()