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