#!/usr/bin/python
# Copyright 2015 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.
import httplib
import logging
import os
import sys
import urllib2
import common
try:
# Ensure the chromite site-package is installed.
from chromite.lib import *
except ImportError:
import subprocess
build_externals_path = os.path.join(
os.path.dirname(os.path.dirname(os.path.realpath(__file__))),
'utils', 'build_externals.py')
subprocess.check_call([build_externals_path, 'chromiterepo'])
# Restart the script so python now finds the autotest site-packages.
sys.exit(os.execv(__file__, sys.argv))
from autotest_lib.server.hosts import moblab_host
from autotest_lib.site_utils import brillo_common
_DEFAULT_STAGE_PATH_TEMPLATE = 'aue2e/%(use)s'
_DEVSERVER_STAGE_URL_TEMPLATE = ('http://%(moblab)s:%(port)s/stage?'
'local_path=%(stage_dir)s&'
'files=%(stage_files)s')
_DEVSERVER_PAYLOAD_URI_TEMPLATE = ('http://%(moblab)s:%(port)s/static/'
'%(stage_path)s')
_STAGED_PAYLOAD_FILENAME = 'update.gz'
_SPEC_GEN_LABEL = 'gen'
_TEST_JOB_NAME = 'brillo_update_test'
_TEST_NAME = 'autoupdate_EndToEndTest'
_DEFAULT_DEVSERVER_PORT = '8080'
# Snippet of code that runs on the Moblab and returns the type of a payload
# file. Result is either 'delta' or 'full', acordingly.
_GET_PAYLOAD_TYPE = """
import update_payload
p = update_payload(open('%(payload_file)s'))
p.Init()
print 'delta' if p.IsDelta() else 'full'
"""
class PayloadStagingError(brillo_common.BrilloTestError):
"""A failure that occurred while staging an update payload."""
class PayloadGenerationError(brillo_common.BrilloTestError):
"""A failure that occurred while generating an update payload."""
def setup_parser(parser):
"""Add parser options.
@param parser: argparse.ArgumentParser of the script.
"""
parser.add_argument('-t', '--target_payload', metavar='SPEC', required=True,
help='Stage a target payload. This can either be a '
'path to a local payload file, or take the form '
'"%s:DST_IMAGE[:SRC_IMAGE]", in which case a '
'new payload will get generated from SRC_IMAGE '
'(if given) and DST_IMAGE and staged on the '
'server. This is a mandatory input.' %
_SPEC_GEN_LABEL)
parser.add_argument('-s', '--source_payload', metavar='SPEC',
help='Stage a source payload. This is an optional '
'input. See --target_payload for possible values '
'for SPEC.')
brillo_common.setup_test_action_parser(parser)
def get_stage_rel_path(stage_file):
"""Returns the relative stage path for remote file.
The relative stage path consists of the last three path components: the
file name and the two directory levels that contain it.
@param stage_file: Path to the file that is being staged.
@return A stage relative path.
"""
components = []
for i in range(3):
stage_file, component = os.path.split(stage_file)
components.insert(0, component)
return os.path.join(*components)
def stage_remote_payload(moblab, devserver_port, tmp_stage_file):
"""Stages a remote payload on the Moblab's devserver.
@param moblab: MoblabHost representing the MobLab being used for testing.
@param devserver_port: Externally accessible port to the Moblab devserver.
@param tmp_stage_file: Path to the remote payload file to stage.
@return URI to use for downloading the staged payload.
@raise PayloadStagingError: If we failed to stage the payload.
"""
# Remove the artifact if previously staged.
stage_rel_path = get_stage_rel_path(tmp_stage_file)
target_stage_file = os.path.join(moblab_host.MOBLAB_IMAGE_STORAGE,
stage_rel_path)
moblab.run('rm -f %s && chown moblab:moblab %s' %
(target_stage_file, tmp_stage_file))
tmp_stage_dir, stage_file = os.path.split(tmp_stage_file)
devserver_host = moblab.web_address.split(':')[0]
try:
stage_url = _DEVSERVER_STAGE_URL_TEMPLATE % {
'moblab': devserver_host,
'port': devserver_port,
'stage_dir': tmp_stage_dir,
'stage_files': stage_file}
res = urllib2.urlopen(stage_url).read()
except (urllib2.HTTPError, httplib.HTTPException, urllib2.URLError) as e:
raise PayloadStagingError('Unable to stage payload on moblab: %s' % e)
else:
if res != 'Success':
raise PayloadStagingError('Staging failed: %s' % res)
logging.debug('Payload is staged on Moblab as %s', stage_rel_path)
return _DEVSERVER_PAYLOAD_URI_TEMPLATE % {
'moblab': devserver_host,
'port': _DEFAULT_DEVSERVER_PORT,
'stage_path': os.path.dirname(stage_rel_path)}
def stage_local_payload(moblab, devserver_port, tmp_stage_dir, payload):
"""Stages a local payload on the MobLab's devserver.
@param moblab: MoblabHost representing the MobLab being used for testing.
@param devserver_port: Externally accessible port to the Moblab devserver.
@param tmp_stage_dir: Path of temporary staging directory on the Moblab.
@param payload: Path to the local payload file to stage.
@return Tuple consisting a payload download URI and the payload type
('delta' or 'full').
@raise PayloadStagingError: If we failed to stage the payload.
"""
if not os.path.isfile(payload):
raise PayloadStagingError('Payload file %s does not exist.' % payload)
# Copy the payload file over to the temporary stage directory.
tmp_stage_file = os.path.join(tmp_stage_dir, _STAGED_PAYLOAD_FILENAME)
moblab.send_file(payload, tmp_stage_file)
# Find the payload type.
get_payload_type = _GET_PAYLOAD_TYPE % {'payload_file': tmp_stage_file}
payload_type = moblab.run('python', stdin=get_payload_type).stdout.strip()
# Stage the copied payload.
payload_uri = stage_remote_payload(moblab, devserver_port, tmp_stage_file)
return payload_uri, payload_type
def generate_payload(moblab, devserver_port, tmp_stage_dir, payload_spec):
"""Generates and stages a payload from local image(s).
@param moblab: MoblabHost representing the MobLab being used for testing.
@param devserver_port: Externally accessible port to the Moblab devserver.
@param tmp_stage_dir: Path of temporary staging directory on the Moblab.
@param payload_spec: A string of the form "DST_IMAGE[:SRC_IMAGE]", where
DST_IMAGE is a target image and SRC_IMAGE an optional
source image.
@return Tuple consisting a payload download URI and the payload type
('delta' or 'full').
@raise PayloadGenerationError: If we failed to generate the payload.
@raise PayloadStagingError: If we failed to stage the payload.
"""
parts = payload_spec.split(':', 1)
dst_image = parts[0]
src_image = parts[1] if len(parts) == 2 else None
if not os.path.isfile(dst_image):
raise PayloadGenerationError('Target image file %s does not exist.' %
dst_image)
if src_image and not os.path.isfile(src_image):
raise PayloadGenerationError('Source image file %s does not exist.' %
src_image)
tmp_images_dir = moblab.make_tmp_dir()
try:
# Copy the images to a temporary location.
remote_dst_image = os.path.join(tmp_images_dir,
os.path.basename(dst_image))
moblab.send_file(dst_image, remote_dst_image)
remote_src_image = None
if src_image:
remote_src_image = os.path.join(tmp_images_dir,
os.path.basename(src_image))
moblab.send_file(src_image, remote_src_image)
# Generate the payload into a temporary staging directory.
tmp_stage_file = os.path.join(tmp_stage_dir, _STAGED_PAYLOAD_FILENAME)
gen_cmd = ['brillo_update_payload', 'generate',
'--payload', tmp_stage_file,
'--target_image', remote_dst_image]
if remote_src_image:
payload_type = 'delta'
gen_cmd += ['--source_image', remote_src_image]
else:
payload_type = 'full'
moblab.run(' '.join(gen_cmd), stdout_tee=None, stderr_tee=None)
finally:
moblab.run('rm -rf %s' % tmp_images_dir)
# Stage the generated payload.
payload_uri = stage_remote_payload(moblab, devserver_port, tmp_stage_file)
return payload_uri, payload_type
def stage_payload(moblab, devserver_port, tmp_dir, use, payload_spec):
"""Stages the payload based on a given specification.
@param moblab: MoblabHost representing the MobLab being used for testing.
@param devserver_port: Externally accessible port to the Moblab devserver.
@param tmp_dir: Path of temporary static subdirectory.
@param use: String defining the use for the payload, either 'source' or
'target'.
@param payload_spec: Either a string of the form
"PAYLOAD:DST_IMAGE[:SRC_IMAGE]" describing how to
generate a new payload from a target and (optionally)
source image; or path to a local payload file.
@return Tuple consisting a payload download URI and the payload type
('delta' or 'full').
@raise PayloadGenerationError: If we failed to generate the payload.
@raise PayloadStagingError: If we failed to stage the payload.
"""
tmp_stage_dir = os.path.join(
tmp_dir, _DEFAULT_STAGE_PATH_TEMPLATE % {'use': use})
moblab.run('mkdir -p %s && chown -R moblab:moblab %s' %
(tmp_stage_dir, tmp_stage_dir))
spec_gen_prefix = _SPEC_GEN_LABEL + ':'
if payload_spec.startswith(spec_gen_prefix):
return generate_payload(moblab, devserver_port, tmp_stage_dir,
payload_spec[len(spec_gen_prefix):])
else:
return stage_local_payload(moblab, devserver_port, tmp_stage_dir,
payload_spec)
def main(args):
"""The main function."""
args = brillo_common.parse_args(
'Set up Moblab for running Brillo AU end-to-end test, then launch '
'the test (unless otherwise requested).',
setup_parser=setup_parser)
moblab, devserver_port = brillo_common.get_moblab_and_devserver_port(
args.moblab_host)
tmp_dir = moblab.make_tmp_dir(base=moblab_host.MOBLAB_IMAGE_STORAGE)
moblab.run('chown -R moblab:moblab %s' % tmp_dir)
test_args = {'name': _TEST_JOB_NAME}
try:
if args.source_payload:
payload_uri, _ = stage_payload(moblab, devserver_port, tmp_dir,
'source', args.source_payload)
test_args['source_payload_uri'] = payload_uri
logging.info('Source payload was staged')
payload_uri, payload_type = stage_payload(
moblab, devserver_port, tmp_dir, 'target', args.target_payload)
test_args['target_payload_uri'] = payload_uri
test_args['update_type'] = payload_type
logging.info('Target payload was staged')
finally:
moblab.run('rm -rf %s' % tmp_dir)
brillo_common.do_test_action(args, moblab, _TEST_NAME, test_args)
if __name__ == '__main__':
try:
main(sys.argv)
sys.exit(0)
except brillo_common.BrilloTestError as e:
logging.error('Error: %s', e)
sys.exit(1)