1#
2# Copyright (C) 2018 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the 'License');
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an 'AS IS' BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16
17import imp  # Python v2 compatibility
18import logging
19import os
20import re
21import subprocess
22import zipfile
23
24from host_controller import common
25from host_controller.command_processor import base_command_processor
26from host_controller.utils.gcp import gcs_utils
27
28from vti.dashboard.proto import TestSuiteResultMessage_pb2 as SuiteResMsg
29from vti.test_serving.proto import TestScheduleConfigMessage_pb2 as SchedCfgMsg
30
31
32class CommandReproduce(base_command_processor.BaseCommandProcessor):
33    """Command processor for reproduce command.
34
35    Attributes:
36        campaign_common: campaign module. Dynamically imported since
37                         the campaign might need to be separated from the
38                         host controller itself.
39    """
40
41    command = "reproduce"
42    command_detail = ("Reproduce the test environment for a pre-run test and "
43                      "execute the tradefed command prompt of the fetched "
44                      "test suite. Setup the test env "
45                      "(fetching, flashing devices, etc) and retrieve "
46                      "formerly run test result to retry on, if the path "
47                      "to the report protobuf file is given.")
48
49    # @Override
50    def SetUp(self):
51        """Initializes the parser for reproduce command."""
52        self.campaign_common = None
53        self.arg_parser.add_argument(
54            "--suite",
55            default="vts",
56            choices=("vts", "cts", "gts", "sts"),
57            help="To specify the type of a test suite to be run.")
58        self.arg_parser.add_argument(
59            "--report_path",
60            required=True,
61            help="Google Cloud Storage URL, the path of a report protobuf file."
62        )
63        self.arg_parser.add_argument(
64            "--serial",
65            default=None,
66            help="The serial numbers for flashing and testing. "
67            "Multiple serial numbers are separated by commas.")
68        self.arg_parser.add_argument(
69            "--automated_retry",
70            action="store_true",
71            help="Retries automatically until all test cases are passed "
72            "or the number or the failed test cases is the same as "
73            "the previous one.")
74
75    # @Override
76    def Run(self, arg_line):
77        """Reproduces the test env of the pre-run test."""
78        args = self.arg_parser.ParseLine(arg_line)
79
80        if args.report_path:
81            gsutil_path = gcs_utils.GetGsutilPath()
82            if not gsutil_path:
83                logging.error(
84                    "Please check whether gsutil is installed and on your PATH"
85                )
86                return False
87
88            if (not args.report_path.startswith("gs://")
89                    or not gcs_utils.IsGcsFile(gsutil_path, args.report_path)):
90                logging.error("%s is not a valid GCS path.", args.report_path)
91                return False
92
93            dest_path = os.path.join("".join(self.ReplaceVars(["{tmp_dir}"])),
94                                     os.path.basename(args.report_path))
95            gcs_utils.Copy(gsutil_path, args.report_path, dest_path)
96            report_msg = SuiteResMsg.TestSuiteResultMessage()
97            try:
98                with open(dest_path, "r") as report_fd:
99                    report_msg.ParseFromString(report_fd.read())
100            except IOError as e:
101                logging.exception(e)
102                return False
103            serial = []
104            if args.serial:
105                serial = args.serial.split(",")
106            setup_command_list = self.GenerateSetupCommands(report_msg, serial)
107            if not setup_command_list:
108                suite_fetch_command = self.GenerateTestSuiteFetchCommand(
109                    report_msg)
110                if suite_fetch_command:
111                    setup_command_list.append(suite_fetch_command)
112            for command in setup_command_list:
113                self.console.onecmd(command)
114
115            if not self.GetResultFromGCS(gsutil_path, report_msg, args.suite):
116                return False
117        else:
118            logging.error("Path to a report protobuf file is required.")
119            return False
120
121        if args.suite not in self.console.test_suite_info:
122            logging.error("test_suite_info doesn't have '%s': %s", args.suite,
123                          self.console.test_suite_info)
124            return False
125
126        if args.automated_retry:
127            if self.campaign_common is None:
128                self.campaign_common = imp.load_source(
129                    'campaign_common',
130                    os.path.join(os.getcwd(), "host_controller", "campaigns",
131                                 "campaign_common.py"))
132            retry_command = self.campaign_common.GenerateRetryCommand(
133                report_msg.schedule_config.build_target[0].name,
134                report_msg.branch, report_msg.suite_name.lower(),
135                report_msg.suite_plan, serial)
136            self.console.onecmd(retry_command)
137        else:
138            subprocess.call(self.console.test_suite_info[args.suite])
139
140    def GenerateSetupCommands(self, report_msg, serial):
141        """Generates fetch, flash commands using fetch info from report_msg.
142
143        Args:
144            report_msg: pb2, contains fetch info of the test suite.
145            serial: list of string, serial number(s) of the device(s)
146                    to be flashed.
147
148        Returns:
149            list of string, console commands to fetch device/gsi images
150            and flash the device(s).
151        """
152        ret = []
153        schedule_config = report_msg.schedule_config
154        if not schedule_config.manifest_branch:
155            logging.error("Report contains no fetch information. "
156                          "Aborting pre-test setups on the device(s).")
157        elif not serial:
158            logging.error("Device serial number(s) not given. "
159                          "Aborting pre-test setups on the device(s).")
160        else:
161            try:
162                build_target_msg = schedule_config.build_target[0]
163                test_schedule_msg = build_target_msg.test_schedule[0]
164            except IndexError as e:
165                logging.exception(e)
166                return ret
167            kwargs = {}
168            # common fetch info
169            kwargs["shards"] = str(len(serial))
170            kwargs["test_name"] = "%s/%s" % (report_msg.suite_name.lower(),
171                                             report_msg.suite_plan)
172            kwargs["serial"] = serial
173
174            # fetch info for device images
175            kwargs["manifest_branch"] = schedule_config.manifest_branch
176            kwargs["build_target"] = build_target_msg.name
177            kwargs["build_id"] = report_msg.vendor_build_id
178            kwargs["pab_account_id"] = schedule_config.pab_account_id
179            if kwargs["manifest_branch"].startswith("gs://"):
180                kwargs[
181                    "build_storage_type"] = SchedCfgMsg.BUILD_STORAGE_TYPE_GCS
182            else:
183                kwargs[
184                    "build_storage_type"] = SchedCfgMsg.BUILD_STORAGE_TYPE_PAB
185            kwargs["require_signed_device_build"] = (
186                build_target_msg.require_signed_device_build)
187            kwargs["has_bootloader_img"] = build_target_msg.has_bootloader_img
188            kwargs["has_radio_img"] = build_target_msg.has_radio_img
189
190            # fetch info for gsi images and gsispl command
191            kwargs["gsi_branch"] = test_schedule_msg.gsi_branch
192            kwargs["gsi_build_target"] = test_schedule_msg.gsi_build_target
193            kwargs["gsi_build_id"] = report_msg.gsi_build_id
194            kwargs["gsi_pab_account_id"] = test_schedule_msg.gsi_pab_account_id
195            if kwargs["gsi_branch"].startswith("gs://"):
196                kwargs["gsi_storage_type"] = SchedCfgMsg.BUILD_STORAGE_TYPE_GCS
197            else:
198                kwargs["gsi_storage_type"] = SchedCfgMsg.BUILD_STORAGE_TYPE_PAB
199            kwargs["gsi_vendor_version"] = test_schedule_msg.gsi_vendor_version
200
201            # fetch info for test suite
202            kwargs["test_build_id"] = report_msg.build_id
203            kwargs["test_branch"] = report_msg.branch
204            kwargs["test_build_target"] = report_msg.target
205            if kwargs["test_build_target"].endswith(".zip"):
206                kwargs["test_build_target"] = kwargs["test_build_target"][:-4]
207            kwargs[
208                "test_pab_account_id"] = test_schedule_msg.test_pab_account_id
209            if kwargs["test_branch"].startswith("gs://"):
210                kwargs[
211                    "test_storage_type"] = SchedCfgMsg.BUILD_STORAGE_TYPE_GCS
212            else:
213                kwargs["gsi_storage_type"] = SchedCfgMsg.BUILD_STORAGE_TYPE_PAB
214
215            self.campaign_common = imp.load_source(
216                "campaign_common",
217                os.path.join(os.getcwd(), "host_controller", "campaigns",
218                             "campaign_common.py"))
219            fetch_commands_result, gsi = self.campaign_common.EmitFetchCommands(
220                **kwargs)
221            ret.extend(fetch_commands_result)
222            flash_commands_result = self.campaign_common.EmitFlashCommands(
223                gsi, **kwargs)
224            ret.extend(flash_commands_result)
225
226        return ret
227
228    def GenerateTestSuiteFetchCommand(self, report_msg):
229        """Generates a fetch command line using fetch info from report_msg.
230
231        Args:
232            report_msg: pb2, contains fetch info of the test suite.
233
234        Returns:
235            string, console command to fetch a test suite artifact.
236        """
237        ret = "fetch"
238
239        if report_msg.branch.startswith("gs://"):
240            ret += " --type=gcs --path=%s/%s --set_suite_as=%s" % (
241                report_msg.branch, report_msg.target,
242                report_msg.suite_name.lower())
243        else:
244            ret += (" --type=pab --branch=%s --target=%s --build_id=%s"
245                    " --artifact_name=android-%s.zip") % (
246                        report_msg.branch, report_msg.target,
247                        report_msg.build_id, report_msg.suite_name.lower())
248            try:
249                build_target_msg = report_msg.schedule_config.build_target[0]
250                test_schedule_msg = build_target_msg.test_schedule[0]
251            except IndexError as e:
252                logging.exception(e)
253                test_schedule_msg = SchedCfgMsg.TestScheduleConfigMessage()
254            if test_schedule_msg.test_pab_account_id:
255                ret += " --account_id=%s" % test_schedule_msg.test_pab_account_id
256            else:
257                ret += " --account_id=%s" % common._DEFAULT_ACCOUNT_ID_INTERNAL
258
259        return ret
260
261    def GetResultFromGCS(self, gsutil_path, report_msg, suite):
262        """Downloads results.zip from GCS and unzip it to the results directory.
263
264        Args:
265            gsutil_path: string, path to the gsutil binary.
266            report_msg: pb2, contains fetch info of the test suite.
267            suite: string, specifies the type of test suite fetched.
268
269        Returns:
270            True if successful. False otherwise
271        """
272        result_base_path = report_msg.result_path
273        try:
274            tools_path = os.path.dirname(self.console.test_suite_info[suite])
275        except KeyError as e:
276            logging.exception(e)
277            return False
278        local_results_path = os.path.join(tools_path,
279                                          common._RESULTS_BASE_PATH)
280
281        if not os.path.exists(local_results_path):
282            os.mkdir(local_results_path)
283
284        result_path_list = gcs_utils.List(gsutil_path, result_base_path)
285        result_zip_url = ""
286        for result_path in result_path_list:
287            if re.match(".*results_.*\.zip", result_path):
288                result_zip_url = result_path
289                break
290
291        if (not result_zip_url or not gcs_utils.Copy(
292                gsutil_path, result_zip_url, local_results_path)):
293            logging.error("Fail to copy from %s.", result_base_path)
294            return False
295
296        result_zip_local_path = os.path.join(local_results_path,
297                                             os.path.basename(result_zip_url))
298        with zipfile.ZipFile(result_zip_local_path, mode="r") as zip_ref:
299            zip_ref.extractall(local_results_path)
300
301        return True
302