Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 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 | |
| 6 | import httplib |
| 7 | import logging |
| 8 | import os |
| 9 | import sys |
| 10 | import urllib2 |
| 11 | |
| 12 | import common |
| 13 | try: |
| 14 | # Ensure the chromite site-package is installed. |
| 15 | from chromite.lib import * |
| 16 | except 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)) |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 24 | from autotest_lib.server.hosts import moblab_host |
| 25 | from 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' |
Gilad Arnold | 2d02a35 | 2015-10-22 13:30:44 -0700 | [diff] [blame] | 38 | _DEFAULT_DEVSERVER_PORT = '8080' |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 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 = """ |
| 43 | import update_payload |
| 44 | p = update_payload(open('%(payload_file)s')) |
| 45 | p.Init() |
| 46 | print 'delta' if p.IsDelta() else 'full' |
| 47 | """ |
| 48 | |
| 49 | |
| 50 | class PayloadStagingError(brillo_common.BrilloTestError): |
| 51 | """A failure that occurred while staging an update payload.""" |
| 52 | |
| 53 | |
| 54 | class PayloadGenerationError(brillo_common.BrilloTestError): |
| 55 | """A failure that occurred while generating an update payload.""" |
| 56 | |
| 57 | |
| 58 | def setup_parser(parser): |
| 59 | """Add parser options. |
| 60 | |
| 61 | @param parser: argparse.ArgumentParser of the script. |
| 62 | """ |
Gilad Arnold | 15cd474 | 2015-11-03 11:58:49 -0800 | [diff] [blame] | 63 | parser.add_argument('-t', '--target_payload', metavar='SPEC', required=True, |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 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 | |
Gilad Arnold | 44629d5 | 2015-11-03 07:23:19 -0800 | [diff] [blame] | 76 | brillo_common.setup_test_action_parser(parser) |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 77 | |
| 78 | |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 79 | def 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 | |
| 96 | def 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) |
Gilad Arnold | 2d02a35 | 2015-10-22 13:30:44 -0700 | [diff] [blame] | 114 | devserver_host = moblab.web_address.split(':')[0] |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 115 | try: |
| 116 | stage_url = _DEVSERVER_STAGE_URL_TEMPLATE % { |
Gilad Arnold | 2d02a35 | 2015-10-22 13:30:44 -0700 | [diff] [blame] | 117 | 'moblab': devserver_host, |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 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 % { |
Gilad Arnold | 2d02a35 | 2015-10-22 13:30:44 -0700 | [diff] [blame] | 130 | 'moblab': devserver_host, |
| 131 | 'port': _DEFAULT_DEVSERVER_PORT, |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 132 | 'stage_path': os.path.dirname(stage_rel_path)} |
| 133 | |
| 134 | |
| 135 | def 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 | |
| 165 | def 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 | |
Gilad Arnold | 5db186b | 2015-11-03 12:00:00 -0800 | [diff] [blame] | 192 | tmp_images_dir = moblab.make_tmp_dir() |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 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 | |
| 225 | def 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 | |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 258 | def 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).', |
Gilad Arnold | 15cd474 | 2015-11-03 11:58:49 -0800 | [diff] [blame] | 263 | setup_parser=setup_parser) |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 264 | |
Gilad Arnold | 5db186b | 2015-11-03 12:00:00 -0800 | [diff] [blame] | 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) |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 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 | |
Gilad Arnold | 44629d5 | 2015-11-03 07:23:19 -0800 | [diff] [blame] | 285 | brillo_common.do_test_action(args, moblab, _TEST_NAME, test_args) |
Gilad Arnold | 8eb72af | 2015-10-20 12:26:20 -0700 | [diff] [blame] | 286 | |
| 287 | |
| 288 | if __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) |