blob: e21edd4ce17e2055473603da80b4dca682a40e2f [file] [log] [blame]
# Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""GenerateTestPayloads 'quickly' creates payloads needed for testing."""
import logging
import optparse
import os
import pickle
import re
import shutil
import sys
import tempfile
import constants
import cros_build_lib as cros_lib
from crostestutils.au_test_harness import cros_au_test_harness
from crostestutils.generate_test_payloads import payload_generation_exception
from crostestutils.lib import dev_server_wrapper
from crostestutils.lib import parallel_test_job
from crostestutils.lib import public_key_manager
from crostestutils.lib import test_helper
class UpdatePayload(object):
"""Defines an update payload."""
def __init__(self, target, base, key=None):
self.base = base = target
self.key = key
def __str__(self):
my_repr =
if self.base:
my_repr = self.base + '->' + my_repr
if self.key:
my_repr = my_repr + '+' + self.key
return my_repr
def __eq__(self, other):
return str(self) == str(other)
def __hash__(self):
return hash(str(self))
class UpdatePayloadGenerator(object):
"""Class responsible for generating update payloads."""
def __init__(self, options):
"""Initializes a generator object from parsed options.
options: Parsed options from main().
""" =
self.base = options.base
self.target_signed = None # Set later when creating the image.
# For vm tests we use the _qemu name for the images. Regardless of vm or
# non vm, these no_vm names are guaranteed to be non-qemu base/target names.
self.base_no_vm = self.base
self.target_no_vm =
# Keys.
self.public_key = options.public_key
self.private_key = options.private_key
# Affect what payloads we create.
self.board = options.board
self.full_suite = options.full_suite
self.payloads = set([])
self.nplus1 = options.nplus1
self.vm = options.vm
def _AddUpdatePayload(self, target, base, key=None):
"""Adds a new required update payload. If base is None, a full payload."""
self.payloads.add(UpdatePayload(target, base, key))
def GenerateImagesForTesting(self):
# All vm testing requires a VM'ized target.
if self.vm: = test_helper.CreateVMImage(, self.board)
if self.full_suite:
if self.public_key:
self.target_signed = + '.signed'
if not os.path.exists(self.target_signed):'Creating a signed image for signed payload test.')
shutil.copy(, self.target_signed)
self.base = test_helper.CreateVMImage(self.base, self.board)
def GeneratePayloadRequirements(self):
"""Generate Payload Requirements for AUTestHarness and NPlus1 Testing."""
def AddPayloadsForAUTestHarness():
if self.full_suite:
self._AddUpdatePayload(self.target_no_vm, self.base)
self._AddUpdatePayload(self.base_no_vm, self.target_no_vm)
# Need a signed payload for the signed payload test.
if self.target_signed:
self._AddUpdatePayload(self.target_no_vm, self.target_signed,
def AddNPlus1Updates():
self._AddUpdatePayload(self.target_no_vm, self.base_no_vm)
self._AddUpdatePayload(self.target_no_vm, self.target_no_vm)
self._AddUpdatePayload(self.target_no_vm, None)
if self.nplus1: AddNPlus1Updates()
def GeneratePayloads(self):
"""Iterates through payload requirements and generates them.
This is the main method of this class. It iterates through payloads
it needs, generates them, and builds a Cache that can be used by the
test harness to reference these payloads.
The cache as a Python dict.
def GeneratePayload(payload, log_file):
"""Returns the error code from generating an update with the devserver."""
# Base command.
command = ['sudo', 'start_devserver', '--pregenerate_update', '--exit']
in_chroot_key = None
in_chroot_base = None
in_chroot_target = cros_lib.ReinterpretPathForChroot(
if payload.base:
in_chroot_base = cros_lib.ReinterpretPathForChroot(payload.base)
if payload.key:
in_chroot_key = cros_lib.ReinterpretPathForChroot(payload.key)
command.append('--image=%s' % in_chroot_target)
if payload.base: command.append('--src_image=%s' % in_chroot_base)
if self.vm: command.append('--for_vm')
if payload.key: command.append('--private_key=%s' % in_chroot_key)
debug_message = 'delta payload from %s to %s' % (payload.base,
if in_chroot_key:
debug_message = 'Generating a signed %s' % debug_message
debug_message = 'Generating an unsigned %s' % debug_message
return cros_lib.RunCommand(command, enter_chroot=True, print_cmd=False,
log_to_file=log_file, error_ok=True,
def ProcessOutput(log_files, return_codes):
"""Processes results from the log files of GeneratePayload invocations.
log_files: A list of filename strings with stored logs.
return_codes: An equally sized list of return codes pertaining to
the processes that ran that generated those logs.
An array of cache entries from the log files.
# Looking for this line in the output.
key_line_re = re.compile('^PREGENERATED_UPDATE=([\w/./+]+)')
return_array = []
for log_file, return_code in zip(log_files, return_codes):
with open(log_file) as log_file_handle:
output =
if return_code != 0:
raise payload_generation_exception.PayloadGenerationException(
'Failed to generate a required update.')
for line in output.splitlines():
match =
if match:
# Convert blah/blah/update.gz -> update/blah/blah.
path_to_update_gz =
path_to_update_dir = path_to_update_gz.rpartition(
# Check that we could actually parse the directory correctly.000
if not path_to_update_dir:
raise payload_generation_exception.PayloadGenerationException(
'Payload generated but failed to parse cache directory.')
return_array.append('/'.join(['update', path_to_update_dir]))
logging.error('Could not find update string in log.')
return return_array
update_ids = []
jobs = []
args = []
log_files = []
for payload in self.payloads:
fd, log_file = tempfile.mkstemp('GenerateVMUpdate')
os.close(fd) # Just want filename so close file immediately.
update_id = dev_server_wrapper.GenerateUpdateId(, payload.base, payload.key)
args.append((payload, log_file))
log_files.append(log_file)'Generating updates required for this test suite in parallel.')
error_codes = parallel_test_job.RunParallelJobs(
test_helper.CalculateDefaultJobs(), jobs, args)
results = ProcessOutput(log_files, error_codes)
# Build the dictionary from our id's and returned cache paths.
cache_dictionary = {}
for index, update_id in enumerate(update_ids):
cache_dictionary[update_id] = results[index]
return cache_dictionary
def DumpCacheToDisk(self, cache):
"""Dumps the cache to the same folder as the images."""
path_to_dump = os.path.dirname(
cache_file = os.path.join(path_to_dump, cros_au_test_harness.CACHE_FILE)'Dumping %s', cache_file)
with open(cache_file, 'w') as file_handle:
pickle.dump(cache, file_handle)
def CheckOptions(parser, options):
"""Checks that given options are valid.
parser: Parser used to parse options.
options: Parse options from OptionParser.
if not or not os.path.isfile(
parser.error('Target image must exist.')
if not options.base:'Base image not specified. Using target as base image.')
options.base =
if not os.path.isfile(options.base):
parser.error('Base image must exist.')
if options.private_key:
if not os.path.isfile(options.private_key):
parser.error('Private key must exist.')
if not os.path.isfile(options.public_key):
parser.error('Public key must exist.')
if options.vm:
if not options.board:
parser.error('Board must be set to generate update '
'payloads for vm.')
def main():
parser = optparse.OptionParser()
parser.add_option('--base', help='Image we want to test updates from.')
parser.add_option('--board', help='Board used for the images.')
parser.add_option('--clean', default=False, dest='clean', action='store_true',
help='Clean cache of previous payloads')
parser.add_option('--full_suite', default=False, action='store_true',
help='Prepare to run the full au test suite.')
parser.add_option('--novm', default=True, action='store_false', dest='vm',
help='Test Harness payloads will not be tested in a VM.')
parser.add_option('--nplus1', default=False, action='store_true',
help='Produce nplus1 updates for testing in lab.')
help='Private key to sign payloads.')
help='Public key to verify signed payloads.')
parser.add_option('--target', help='Image we want to test updates to.')
options = parser.parse_args()[0]
CheckOptions(parser, options)
if options.clean:
generator = UpdatePayloadGenerator(options)
cache = generator.GeneratePayloads()
if __name__ == '__main__':