1# Copyright (c) 2014 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.
4import logging
5import re
6
7from autotest_lib.client.common_lib import error, global_config
8from autotest_lib.server import test
9from autotest_lib.server.hosts import moblab_host
10
11
12DEFAULT_IMAGE_STORAGE_SERVER = global_config.global_config.get_config_value(
13        'CROS', 'image_storage_server')
14STORAGE_SERVER_REGEX = 'gs://.*/'
15
16
17class MoblabTest(test.test):
18    """Base class for Moblab tests.
19    """
20
21    def initialize(self, host, boto_path='',
22                   image_storage_server=DEFAULT_IMAGE_STORAGE_SERVER):
23        """Initialize the Moblab Host.
24
25        * Installs a boto file.
26        * Sets up the image storage server for this test.
27        * Finds and adds DUTs on the testing subnet.
28
29        @param boto_path: Path to the boto file we want to install.
30        @param image_storage_server: image storage server to use for grabbing
31                                     images from Google Storage.
32        """
33        super(MoblabTest, self).initialize()
34        self._host = host
35        self._host.install_boto_file(boto_path)
36        self.set_image_storage_server(image_storage_server)
37        self._host.wait_afe_up()
38        self._host.find_and_add_duts()
39
40
41    def set_image_storage_server(self, image_storage_server):
42        """Set the image storage server.
43
44        @param image_storage_server: Name of image storage server to use. Must
45                                     follow format or gs://bucket-name/
46                                     (Note trailing slash is required).
47
48        @raises error.TestError if the image_storage_server is incorrectly
49                                formatted.
50        """
51        if not re.match(STORAGE_SERVER_REGEX, image_storage_server):
52            raise error.TestError(
53                    'Image Storage Server supplied is not in the correct '
54                    'format. Must start with gs:// and end with a trailing '
55                    'slash: %s' % image_storage_server)
56        logging.debug('Setting image_storage_server to %s',
57                      image_storage_server)
58        # If the image_storage_server is already set, delete it.
59        self._host.run('sed -i /image_storage_server/d %s' %
60                       moblab_host.SHADOW_CONFIG_PATH, ignore_status=True)
61        self._host.run("sed -i '/\[CROS\]/ a\image_storage_server: "
62                       "%s' %s" %(image_storage_server,
63                                  moblab_host.SHADOW_CONFIG_PATH))