1#!/usr/bin/python 2# Copyright 2015 The Chromium OS Authors. All rights reserved. 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6import httplib 7import logging 8import os 9import sys 10import urllib2 11 12import common 13try: 14 # Ensure the chromite site-package is installed. 15 from chromite.lib import * 16except ImportError: 17 import subprocess 18 build_externals_path = os.path.join( 19 os.path.dirname(os.path.dirname(os.path.realpath(__file__))), 20 'utils', 'build_externals.py') 21 subprocess.check_call([build_externals_path, 'chromiterepo']) 22 # Restart the script so python now finds the autotest site-packages. 23 sys.exit(os.execv(__file__, sys.argv)) 24from autotest_lib.server.hosts import moblab_host 25from autotest_lib.site_utils import brillo_common 26 27 28_DEFAULT_STAGE_PATH_TEMPLATE = 'aue2e/%(use)s' 29_DEVSERVER_STAGE_URL_TEMPLATE = ('http://%(moblab)s:%(port)s/stage?' 30 'local_path=%(stage_dir)s&' 31 'files=%(stage_files)s') 32_DEVSERVER_PAYLOAD_URI_TEMPLATE = ('http://%(moblab)s:%(port)s/static/' 33 '%(stage_path)s') 34_STAGED_PAYLOAD_FILENAME = 'update.gz' 35_SPEC_GEN_LABEL = 'gen' 36_TEST_JOB_NAME = 'brillo_update_test' 37_TEST_NAME = 'autoupdate_EndToEndTest' 38_DEFAULT_DEVSERVER_PORT = '8080' 39 40# Snippet of code that runs on the Moblab and returns the type of a payload 41# file. Result is either 'delta' or 'full', acordingly. 42_GET_PAYLOAD_TYPE = """ 43import update_payload 44p = update_payload(open('%(payload_file)s')) 45p.Init() 46print 'delta' if p.IsDelta() else 'full' 47""" 48 49 50class PayloadStagingError(brillo_common.BrilloTestError): 51 """A failure that occurred while staging an update payload.""" 52 53 54class PayloadGenerationError(brillo_common.BrilloTestError): 55 """A failure that occurred while generating an update payload.""" 56 57 58def setup_parser(parser): 59 """Add parser options. 60 61 @param parser: argparse.ArgumentParser of the script. 62 """ 63 parser.add_argument('-t', '--target_payload', metavar='SPEC', required=True, 64 help='Stage a target payload. This can either be a ' 65 'path to a local payload file, or take the form ' 66 '"%s:DST_IMAGE[:SRC_IMAGE]", in which case a ' 67 'new payload will get generated from SRC_IMAGE ' 68 '(if given) and DST_IMAGE and staged on the ' 69 'server. This is a mandatory input.' % 70 _SPEC_GEN_LABEL) 71 parser.add_argument('-s', '--source_payload', metavar='SPEC', 72 help='Stage a source payload. This is an optional ' 73 'input. See --target_payload for possible values ' 74 'for SPEC.') 75 76 brillo_common.setup_test_action_parser(parser) 77 78 79def get_stage_rel_path(stage_file): 80 """Returns the relative stage path for remote file. 81 82 The relative stage path consists of the last three path components: the 83 file name and the two directory levels that contain it. 84 85 @param stage_file: Path to the file that is being staged. 86 87 @return A stage relative path. 88 """ 89 components = [] 90 for i in range(3): 91 stage_file, component = os.path.split(stage_file) 92 components.insert(0, component) 93 return os.path.join(*components) 94 95 96def stage_remote_payload(moblab, devserver_port, tmp_stage_file): 97 """Stages a remote payload on the Moblab's devserver. 98 99 @param moblab: MoblabHost representing the MobLab being used for testing. 100 @param devserver_port: Externally accessible port to the Moblab devserver. 101 @param tmp_stage_file: Path to the remote payload file to stage. 102 103 @return URI to use for downloading the staged payload. 104 105 @raise PayloadStagingError: If we failed to stage the payload. 106 """ 107 # Remove the artifact if previously staged. 108 stage_rel_path = get_stage_rel_path(tmp_stage_file) 109 target_stage_file = os.path.join(moblab_host.MOBLAB_IMAGE_STORAGE, 110 stage_rel_path) 111 moblab.run('rm -f %s && chown moblab:moblab %s' % 112 (target_stage_file, tmp_stage_file)) 113 tmp_stage_dir, stage_file = os.path.split(tmp_stage_file) 114 devserver_host = moblab.web_address.split(':')[0] 115 try: 116 stage_url = _DEVSERVER_STAGE_URL_TEMPLATE % { 117 'moblab': devserver_host, 118 'port': devserver_port, 119 'stage_dir': tmp_stage_dir, 120 'stage_files': stage_file} 121 res = urllib2.urlopen(stage_url).read() 122 except (urllib2.HTTPError, httplib.HTTPException, urllib2.URLError) as e: 123 raise PayloadStagingError('Unable to stage payload on moblab: %s' % e) 124 else: 125 if res != 'Success': 126 raise PayloadStagingError('Staging failed: %s' % res) 127 128 logging.debug('Payload is staged on Moblab as %s', stage_rel_path) 129 return _DEVSERVER_PAYLOAD_URI_TEMPLATE % { 130 'moblab': devserver_host, 131 'port': _DEFAULT_DEVSERVER_PORT, 132 'stage_path': os.path.dirname(stage_rel_path)} 133 134 135def stage_local_payload(moblab, devserver_port, tmp_stage_dir, payload): 136 """Stages a local payload on the MobLab's devserver. 137 138 @param moblab: MoblabHost representing the MobLab being used for testing. 139 @param devserver_port: Externally accessible port to the Moblab devserver. 140 @param tmp_stage_dir: Path of temporary staging directory on the Moblab. 141 @param payload: Path to the local payload file to stage. 142 143 @return Tuple consisting a payload download URI and the payload type 144 ('delta' or 'full'). 145 146 @raise PayloadStagingError: If we failed to stage the payload. 147 """ 148 if not os.path.isfile(payload): 149 raise PayloadStagingError('Payload file %s does not exist.' % payload) 150 151 # Copy the payload file over to the temporary stage directory. 152 tmp_stage_file = os.path.join(tmp_stage_dir, _STAGED_PAYLOAD_FILENAME) 153 moblab.send_file(payload, tmp_stage_file) 154 155 # Find the payload type. 156 get_payload_type = _GET_PAYLOAD_TYPE % {'payload_file': tmp_stage_file} 157 payload_type = moblab.run('python', stdin=get_payload_type).stdout.strip() 158 159 # Stage the copied payload. 160 payload_uri = stage_remote_payload(moblab, devserver_port, tmp_stage_file) 161 162 return payload_uri, payload_type 163 164 165def generate_payload(moblab, devserver_port, tmp_stage_dir, payload_spec): 166 """Generates and stages a payload from local image(s). 167 168 @param moblab: MoblabHost representing the MobLab being used for testing. 169 @param devserver_port: Externally accessible port to the Moblab devserver. 170 @param tmp_stage_dir: Path of temporary staging directory on the Moblab. 171 @param payload_spec: A string of the form "DST_IMAGE[:SRC_IMAGE]", where 172 DST_IMAGE is a target image and SRC_IMAGE an optional 173 source image. 174 175 @return Tuple consisting a payload download URI and the payload type 176 ('delta' or 'full'). 177 178 @raise PayloadGenerationError: If we failed to generate the payload. 179 @raise PayloadStagingError: If we failed to stage the payload. 180 """ 181 parts = payload_spec.split(':', 1) 182 dst_image = parts[0] 183 src_image = parts[1] if len(parts) == 2 else None 184 185 if not os.path.isfile(dst_image): 186 raise PayloadGenerationError('Target image file %s does not exist.' % 187 dst_image) 188 if src_image and not os.path.isfile(src_image): 189 raise PayloadGenerationError('Source image file %s does not exist.' % 190 src_image) 191 192 tmp_images_dir = moblab.make_tmp_dir() 193 try: 194 # Copy the images to a temporary location. 195 remote_dst_image = os.path.join(tmp_images_dir, 196 os.path.basename(dst_image)) 197 moblab.send_file(dst_image, remote_dst_image) 198 remote_src_image = None 199 if src_image: 200 remote_src_image = os.path.join(tmp_images_dir, 201 os.path.basename(src_image)) 202 moblab.send_file(src_image, remote_src_image) 203 204 # Generate the payload into a temporary staging directory. 205 tmp_stage_file = os.path.join(tmp_stage_dir, _STAGED_PAYLOAD_FILENAME) 206 gen_cmd = ['brillo_update_payload', 'generate', 207 '--payload', tmp_stage_file, 208 '--target_image', remote_dst_image] 209 if remote_src_image: 210 payload_type = 'delta' 211 gen_cmd += ['--source_image', remote_src_image] 212 else: 213 payload_type = 'full' 214 215 moblab.run(' '.join(gen_cmd), stdout_tee=None, stderr_tee=None) 216 finally: 217 moblab.run('rm -rf %s' % tmp_images_dir) 218 219 # Stage the generated payload. 220 payload_uri = stage_remote_payload(moblab, devserver_port, tmp_stage_file) 221 222 return payload_uri, payload_type 223 224 225def stage_payload(moblab, devserver_port, tmp_dir, use, payload_spec): 226 """Stages the payload based on a given specification. 227 228 @param moblab: MoblabHost representing the MobLab being used for testing. 229 @param devserver_port: Externally accessible port to the Moblab devserver. 230 @param tmp_dir: Path of temporary static subdirectory. 231 @param use: String defining the use for the payload, either 'source' or 232 'target'. 233 @param payload_spec: Either a string of the form 234 "PAYLOAD:DST_IMAGE[:SRC_IMAGE]" describing how to 235 generate a new payload from a target and (optionally) 236 source image; or path to a local payload file. 237 238 @return Tuple consisting a payload download URI and the payload type 239 ('delta' or 'full'). 240 241 @raise PayloadGenerationError: If we failed to generate the payload. 242 @raise PayloadStagingError: If we failed to stage the payload. 243 """ 244 tmp_stage_dir = os.path.join( 245 tmp_dir, _DEFAULT_STAGE_PATH_TEMPLATE % {'use': use}) 246 moblab.run('mkdir -p %s && chown -R moblab:moblab %s' % 247 (tmp_stage_dir, tmp_stage_dir)) 248 249 spec_gen_prefix = _SPEC_GEN_LABEL + ':' 250 if payload_spec.startswith(spec_gen_prefix): 251 return generate_payload(moblab, devserver_port, tmp_stage_dir, 252 payload_spec[len(spec_gen_prefix):]) 253 else: 254 return stage_local_payload(moblab, devserver_port, tmp_stage_dir, 255 payload_spec) 256 257 258def main(args): 259 """The main function.""" 260 args = brillo_common.parse_args( 261 'Set up Moblab for running Brillo AU end-to-end test, then launch ' 262 'the test (unless otherwise requested).', 263 setup_parser=setup_parser) 264 265 moblab, devserver_port = brillo_common.get_moblab_and_devserver_port( 266 args.moblab_host) 267 tmp_dir = moblab.make_tmp_dir(base=moblab_host.MOBLAB_IMAGE_STORAGE) 268 moblab.run('chown -R moblab:moblab %s' % tmp_dir) 269 test_args = {'name': _TEST_JOB_NAME} 270 try: 271 if args.source_payload: 272 payload_uri, _ = stage_payload(moblab, devserver_port, tmp_dir, 273 'source', args.source_payload) 274 test_args['source_payload_uri'] = payload_uri 275 logging.info('Source payload was staged') 276 277 payload_uri, payload_type = stage_payload( 278 moblab, devserver_port, tmp_dir, 'target', args.target_payload) 279 test_args['target_payload_uri'] = payload_uri 280 test_args['update_type'] = payload_type 281 logging.info('Target payload was staged') 282 finally: 283 moblab.run('rm -rf %s' % tmp_dir) 284 285 brillo_common.do_test_action(args, moblab, _TEST_NAME, test_args) 286 287 288if __name__ == '__main__': 289 try: 290 main(sys.argv) 291 sys.exit(0) 292 except brillo_common.BrilloTestError as e: 293 logging.error('Error: %s', e) 294 295 sys.exit(1) 296