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 416 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='')
     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_missing_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')
+    self.mock(bot_auth, 'validate_bot_id_and_fetch_config',
+              lambda *args, **kwargs: cfg)
+    self.do_handshake(status=503)
+
+  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')
+    self.mock(bot_auth, 'validate_bot_id_and_fetch_config',
+              lambda *args, **kwargs: cfg)
+    # Inject the script.
+    def get_self_config(name, store_last_good):
+      self.assertEqual(True, store_last_good)
+      if name == 'scripts/bot_config.py':
+        return None, None
+      self.assertEqual('scripts/foo.py', name)
+      return '1', 'import sys; sys.exit(1)'
+    self.mock(handlers_bot.config.config, 'get_self_config', get_self_config)
+    params = self.do_handshake()
+    self.assertEqual(u'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 637 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()