1# Copyright (c) 2013 The Chromium OS Authors. All rights reserved. 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5import json 6import math 7import os 8 9from autotest_lib.server import test 10from autotest_lib.server import utils 11 12VIDEO_LIST = '__test_video_list' 13 14 15class video_VDAStressSetup(test.test): 16 """ 17 Setup for VDA stress test's server by coping a list of videos from a gs 18 bucket to the server. 19 """ 20 version = 1 21 22 def run_once(self, gs_bucket, server_videos_dir, videos, shard_number, 23 shard_count): 24 if not gs_bucket.endswith('/'): 25 gs_bucket += '/' 26 27 # Probably should not use os.path.join for gs:// paths. 28 gs_video_list = '%s%s' % (gs_bucket, VIDEO_LIST) 29 local_video_list = os.path.join(server_videos_dir, VIDEO_LIST) 30 try: 31 utils.system('gsutil cp %s %s' % (gs_video_list, local_video_list)) 32 videos.extend(json.load(open(local_video_list))) 33 finally: 34 os.remove(local_video_list) 35 36 # Break test_video_list into equal sized shards numbered 0 and only 37 # download shard_number. 38 video_count = len(videos) 39 shard_size = int(math.ceil(video_count / float(shard_count))) 40 begin = shard_size * shard_number 41 end = min(video_count, shard_size * (shard_number + 1)) 42 # Enforce sorting even if VIDEO_LIST file is not sorted. 43 videos.sort() 44 videos = videos[begin:end] 45 46 for video in videos: 47 file_name, _, _ = video.partition(':') 48 utils.system('gsutil cp %s %s' % 49 ('%s%s' % (gs_bucket, file_name), server_videos_dir)) 50