Move url_open with dependencies to utils.net module.

swarming.py

 #!/usr/bin/env python
 # Copyright 2013 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Client tool to trigger tasks or retrieve results from a Swarming server."""
 
 __version__ = '0.1'
 
+import binascii
 import hashlib
 import json
 import logging
 import os
 import re
 import shutil
-import StringIO
 import subprocess
 import sys
 import time
 import urllib
-import zipfile
 
 from third_party import colorama
 from third_party.depot_tools import fix_encoding
 from third_party.depot_tools import subcommand
+
+from utils import net
+from utils import threading_utils
 from utils import tools
-from utils import threading_utils
+from utils import zip_package
 
 import run_isolated
 
 
 ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
 TOOLS_PATH = os.path.join(ROOT_DIR, 'tools')
 
 
 # Default servers.
 # TODO(maruel): Chromium-specific.
@@ skipping 38 matching lines @@
         shards - The number of swarm shards to request.
         test_filter - The gtest filter to apply when running the test.
         slave_os - OS to run on.
         working_dir - Relative working directory to start the script.
         isolate_server - isolate server url.
         verbose - if True, have the slave print more details.
         profile - if True, have the slave print more timing data.
         priority - int between 0 and 1000, lower the higher priority
     """
     self.manifest_hash = manifest_hash
+    self.bundle = zip_package.ZipPackage(ROOT_DIR)
+
     self._test_name = test_name
     self._shards = shards
     self._test_filter = test_filter
     self._target_platform = PLATFORM_MAPPING[slave_os]
     self._working_dir = working_dir
 
     self.data_server_retrieval = isolate_server + '/content/retrieve/default/'
     self._data_server_storage = isolate_server + '/content/store/default/'
     self._data_server_has = isolate_server + '/content/contains/default'
     self._data_server_get_token = isolate_server + '/content/get_token'
 
     self.verbose = bool(verbose)
     self.profile = bool(profile)
     self.priority = priority
 
     self._zip_file_hash = ''
     self._tasks = []
     self._files = {}
     self._token_cache = None
 
   def _token(self):
     if not self._token_cache:
-      result = run_isolated.url_open(self._data_server_get_token)
+      result = net.url_open(self._data_server_get_token)
       if not result:
         # TODO(maruel): Implement authentication.
         raise Failure('Failed to get token, need authentication')
       # Quote it right away, so creating the urls is simpler.
       self._token_cache = urllib.quote(result.read())
     return self._token_cache
 
   def add_task(self, task_name, actions, time_out=600):
     """Appends a new task to the swarm manifest file."""
     # See swarming/src/common/test_request_message.py TestObject constructor for
     # the valid flags.
     self._tasks.append(
         {
           'action': actions,
           'decorate_output': self.verbose,
           'test_name': task_name,
           'time_out': time_out,
         })
 
-  def add_file(self, source_path, rel_path):
-    self._files[source_path] = rel_path
-
   def zip_and_upload(self):
     """Zips up all the files necessary to run a shard and uploads to Swarming
     master.
     """
     assert not self._zip_file_hash
+
     start_time = time.time()
-
-    zip_memory_file = StringIO.StringIO()
-    zip_file = zipfile.ZipFile(zip_memory_file, 'w')
-
-    for source, relpath in self._files.iteritems():
-      zip_file.write(source, relpath)
-
-    zip_file.close()
+    zip_contents = self.bundle.zip_into_buffer()
+    self._zip_file_hash = hashlib.sha1(zip_contents).hexdigest()
     print 'Zipping completed, time elapsed: %f' % (time.time() - start_time)
 
-    zip_memory_file.flush()
-    zip_contents = zip_memory_file.getvalue()
-    zip_memory_file.close()
-
-    self._zip_file_hash = hashlib.sha1(zip_contents).hexdigest()
-
-    response = run_isolated.url_open(
+    response = net.url_open(
         self._data_server_has + '?token=%s' % self._token(),
-        data=self._zip_file_hash,
+        data=binascii.unhexlify(self._zip_file_hash),

[M-A Ruel, 2013/08/28 15:01:48]

You could use .digest() instead.

         content_type='application/octet-stream')
     if response is None:
       print >> sys.stderr, (
           'Unable to query server for zip file presence, aborting.')
       return False
 
     if response.read(1) == chr(1):
       print 'Zip file already on server, no need to reupload.'
       return True
 
     print 'Zip file not on server, starting uploading.'
 
     url = '%s%s?priority=0&token=%s' % (
         self._data_server_storage, self._zip_file_hash, self._token())
-    response = run_isolated.url_open(
+    response = net.url_open(
         url, data=zip_contents, content_type='application/octet-stream')
     if response is None:
       print >> sys.stderr, 'Failed to upload the zip file: %s' % url
       return False
 
     return True
 
   def to_json(self):
     """Exports the current configuration into a swarm-readable manifest file.
 
@@ skipping 35 matching lines @@
 def now():
   """Exists so it can be mocked easily."""
   return time.time()
 
 
 def get_test_keys(swarm_base_url, test_name):
   """Returns the Swarm test key for each shards of test_name."""
   key_data = urllib.urlencode([('name', test_name)])
   url = '%s/get_matching_test_cases?%s' % (swarm_base_url, key_data)
 
-  for i in range(run_isolated.URL_OPEN_MAX_ATTEMPTS):
-    response = run_isolated.url_open(url, retry_404=True)
+  for i in range(net.URL_OPEN_MAX_ATTEMPTS):
+    response = net.url_open(url, retry_404=True)
     if response is None:
       raise Failure(
           'Error: Unable to find any tests with the name, %s, on swarm server'
           % test_name)
 
     result = response.read()
     # TODO(maruel): Compare exact string.
     if 'No matching' in result:
       logging.warning('Unable to find any tests with the name, %s, on swarm '
                       'server' % test_name)
-      if i != run_isolated.URL_OPEN_MAX_ATTEMPTS:
-        run_isolated.HttpService.sleep_before_retry(i, None)
+      if i != net.URL_OPEN_MAX_ATTEMPTS:
+        net.HttpService.sleep_before_retry(i, None)
       continue
     return json.loads(result)
 
   raise Failure(
       'Error: Unable to find any tests with the name, %s, on swarm server'
       % test_name)
 
 
 def retrieve_results(base_url, test_key, timeout, should_stop):
   """Retrieves results for a single test_key."""
   assert isinstance(timeout, float)
   params = [('r', test_key)]
   result_url = '%s/get_result?%s' % (base_url, urllib.urlencode(params))
   start = now()
   while True:
     if timeout and (now() - start) >= timeout:
       logging.error('retrieve_results(%s) timed out', base_url)
       return {}
     # Do retries ourselves.
-    response = run_isolated.url_open(
-        result_url, retry_404=False, retry_50x=False)
+    response = net.url_open(result_url, retry_404=False, retry_50x=False)
     if response is None:
       # Aggressively poll for results. Do not use retry_404 so
       # should_stop is polled more often.
       remaining = min(5, timeout - (now() - start)) if timeout else 5
       if remaining > 0:
-        run_isolated.HttpService.sleep_before_retry(1, remaining)
+        net.HttpService.sleep_before_retry(1, remaining)
     else:
       try:
         data = json.load(response) or {}
       except (ValueError, TypeError):
         logging.warning(
             'Received corrupted data for test_key %s. Retrying.', test_key)
       else:
         if data['output']:
           return data
     if should_stop.get():
@@ skipping 37 matching lines @@
     finally:
       # Done, kill the remaining threads.
       should_stop.set()
 
 
 def chromium_setup(manifest):
   """Sets up the commands to run.
 
   Highly chromium specific.
   """
+  # Add uncompressed zip here. It'll be compressed as part of the package sent
+  # to Swarming server.
+  run_test_name = 'run_isolated.zip'
+  manifest.bundle.add_buffer(run_test_name,
+    run_isolated.get_as_zip_package().zip_into_buffer(compress=False))
+
   cleanup_script_name = 'swarm_cleanup.py'
-  cleanup_script_path = os.path.join(TOOLS_PATH, cleanup_script_name)
-  run_test_name = 'run_isolated.py'
-  run_test_path = os.path.join(ROOT_DIR, run_test_name)
+  manifest.bundle.add_file(os.path.join(TOOLS_PATH, cleanup_script_name),
+    cleanup_script_name)
 
-  manifest.add_file(run_test_path, run_test_name)
-  manifest.add_file(cleanup_script_path, cleanup_script_name)
   run_cmd = [
     'python', run_test_name,
     '--hash', manifest.manifest_hash,
     '--remote', manifest.data_server_retrieval.rstrip('/') + '-gzip/',
   ]
   if manifest.verbose or manifest.profile:
     # Have it print the profiling section.
     run_cmd.append('--verbose')
   manifest.add_task('Run Test', run_cmd)
 
@@ skipping 56 matching lines @@
   print('Zipping up files...')
   if not manifest.zip_and_upload():
     return 1
 
   # Send test requests off to swarm.
   print('Sending test requests to swarm.')
   print('Server: %s' % swarming)
   print('Job name: %s' % test_name)
   test_url = swarming + '/test'
   manifest_text = manifest.to_json()
-  result = run_isolated.url_open(test_url, data={'request': manifest_text})
+  result = net.url_open(test_url, data={'request': manifest_text})
   if not result:
     print >> sys.stderr, 'Failed to send test for %s\n%s' % (
         test_name, test_url)
     return 1
   try:
     json.load(result)
   except (ValueError, TypeError) as e:
     print >> sys.stderr, 'Failed to send test for %s' % test_name
     print >> sys.stderr, 'Manifest: %s' % manifest_text
     print >> sys.stderr, str(e)
@@ skipping 261 matching lines @@
     sys.stderr.write(str(e))
     sys.stderr.write('\n')
     return 1
 
 
 if __name__ == '__main__':
   fix_encoding.fix_encoding()
   tools.disable_buffering()
   colorama.init()
   sys.exit(main(sys.argv[1:]))