server/site_tests/video_VDAStressSetup/video_VDAStressSetup.py

# Copyright (c) 2013 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 json
import math
import os

from autotest_lib.server import test
from autotest_lib.server import utils

VIDEO_LIST = '__test_video_list'


class video_VDAStressSetup(test.test):
    """
    Setup for VDA stress test's server by coping a list of videos from a gs
    bucket to the server.
    """
    version = 1

    def run_once(self, gs_bucket, server_videos_dir, videos, shard_number,
                 shard_count):
        if not gs_bucket.endswith('/'):
            gs_bucket += '/'

        # Probably should not use os.path.join for gs:// paths.
        gs_video_list = '%s%s' % (gs_bucket, VIDEO_LIST)
        local_video_list = os.path.join(server_videos_dir, VIDEO_LIST)
        try:
            utils.system('gsutil cp %s %s' % (gs_video_list, local_video_list))
            videos.extend(json.load(open(local_video_list)))
        finally:
            os.remove(local_video_list)

        # Break test_video_list into equal sized shards numbered 0 and only
        # download shard_number.
        video_count = len(videos)
        shard_size = int(math.ceil(video_count / float(shard_count)))
        begin = shard_size * shard_number
        end = min(video_count, shard_size * (shard_number + 1))
        # Enforce sorting even if VIDEO_LIST file is not sorted.
        videos.sort()
        videos = videos[begin:end]

        for video in videos:
            file_name, _, _ = video.partition(':')
            utils.system('gsutil cp %s %s' %
                         ('%s%s' % (gs_bucket, file_name), server_videos_dir))