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 csv
18import os
19import logging
20import re
21import shutil
22import tempfile
23import zipfile
24
25try:
26    # TODO: Remove when we stop supporting Python 2
27    import StringIO as string_io_module
28except ImportError:
29    import io as string_io_module
30
31import gspread
32
33from oauth2client.service_account import ServiceAccountCredentials
34
35from host_controller import common
36from host_controller.command_processor import base_command_processor
37from host_controller.utils.gcp import gcs_utils
38from host_controller.utils.parser import result_utils
39from host_controller.utils.parser import xml_utils
40
41# Attributes shown on spreadsheet
42_RESULT_ATTR_KEYS = [
43    common._SUITE_NAME_ATTR_KEY, common._SUITE_PLAN_ATTR_KEY,
44    common._SUITE_VERSION_ATTR_KEY, common._SUITE_BUILD_NUM_ATTR_KEY,
45    common._START_DISPLAY_TIME_ATTR_KEY,
46    common._END_DISPLAY_TIME_ATTR_KEY
47]
48
49_BUILD_ATTR_KEYS = [
50    common._FINGERPRINT_ATTR_KEY,
51    common._SYSTEM_FINGERPRINT_ATTR_KEY,
52    common._VENDOR_FINGERPRINT_ATTR_KEY
53]
54
55_SUMMARY_ATTR_KEYS = [
56    common._PASSED_ATTR_KEY, common._FAILED_ATTR_KEY,
57    common._MODULES_TOTAL_ATTR_KEY, common._MODULES_DONE_ATTR_KEY
58]
59
60# Texts on spreadsheet
61_TABLE_HEADER = ("BITNESS", "TEST_MODULE", "TEST_CLASS", "TEST_CASE", "RESULT")
62
63_CMP_TABLE_HEADER = _TABLE_HEADER + ("REFERENCE_RESULT",)
64
65_TOO_MANY_DATA = "too many to be displayed"
66
67
68class CommandSheet(base_command_processor.BaseCommandProcessor):
69    """Command processor for sheet command.
70
71    Attributes:
72        _SCOPE: The scope needed to access Google Sheets.
73        arg_parser: ConsoleArgumentParser object, argument parser.
74        console: cmd.Cmd console object.
75        command: string, command name which this processor will handle.
76        command_detail: string, detailed explanation for the command.
77    """
78    _SCOPE = "https://www.googleapis.com/auth/drive"
79    command = "sheet"
80    command_detail = "Convert and upload a file to Google Sheets."
81
82    # @Override
83    def SetUp(self):
84        """Initializes the parser for sheet command."""
85        self.arg_parser.add_argument(
86            "--src",
87            required=True,
88            help="The local file or GCS URL to be uploaded to Google Sheets. "
89            "This command supports the results produced by TradeFed in XML "
90            "and ZIP formats. Variables enclosed in {} are replaced with the "
91            "the values stored in the console.")
92        self.arg_parser.add_argument(
93            "--dest",
94            required=True,
95            help="The ID of the spreadsheet to which the file is uploaded.")
96        self.arg_parser.add_argument(
97            "--ref",
98            default=None,
99            help="The reference file to be compared with src. If a test in "
100            "src fails, its result in ref is also written to the spreadsheet.")
101        self.arg_parser.add_argument(
102            "--extra_rows",
103            nargs="*",
104            default=[],
105            help="The extra rows written to the spreadsheet. Each argument "
106            "is a row. Cells in a row are separated by commas. Each cell can "
107            "contain variables enclosed in {}.")
108        self.arg_parser.add_argument(
109            "--max",
110            default=30000,
111            type=int,
112            help="Maximum number of results written to the spreadsheet. "
113            "If there are too many results, only failing ones are written.")
114        self.arg_parser.add_argument(
115            "--primary_abi_only",
116            action="store_true",
117            help="Whether to upload only the test results for primary ABI. If "
118            "ref is also specified, this command loads the primary ABI "
119            "results from ref and compares regardless of bitness.")
120        self.arg_parser.add_argument(
121            "--client_secrets",
122            default=None,
123            help="The path to the client secrets file in JSON format for "
124            "authentication. If this argument is not specified, this command "
125            "uses PAB client secrets.")
126
127    # @Override
128    def Run(self, arg_line):
129        """Uploads args.src file to args.dest on Google Sheets."""
130        args = self.arg_parser.ParseLine(arg_line)
131
132        try:
133            src_path = self.console.FormatString(args.src)
134            ref_path = (None if args.ref is None else
135                        self.console.FormatString(args.ref))
136            extra_rows = []
137            for row in args.extra_rows:
138                extra_rows.append([self.console.FormatString(cell)
139                                   for cell in row.split(",")])
140        except KeyError as e:
141            logging.error(
142                "Unknown or uninitialized variable in arguments: %s", e)
143            return False
144
145        if args.client_secrets is not None:
146            credentials = ServiceAccountCredentials.from_json_keyfile_name(
147                args.client_secrets, scopes=self._SCOPE)
148        else:
149            credentials = self.console.build_provider["pab"].Authenticate(
150                scopes=self._SCOPE)
151        client = gspread.authorize(credentials)
152
153        # Load result_attrs, build_attrs, summary_attrs,
154        # src_dict, ref_dict, and exceed_max
155        temp_dir = tempfile.mkdtemp()
156        try:
157            src_path = _GetResultAsXml(src_path, os.path.join(temp_dir, "src"))
158            if not src_path:
159                return False
160
161            with open(src_path, "r") as src_file:
162                (result_attrs,
163                 build_attrs,
164                 summary_attrs) = result_utils.LoadTestSummary(src_file)
165                src_file.seek(0)
166                if args.primary_abi_only:
167                    abis = build_attrs.get(
168                        common._ABIS_ATTR_KEY, "").split(",")
169                    src_bitness = str(result_utils.GetAbiBitness(abis[0]))
170                    src_dict, exceed_max = _LoadSrcResults(src_file, args.max,
171                                                           src_bitness)
172                else:
173                    src_dict, exceed_max = _LoadSrcResults(src_file, args.max)
174
175            if ref_path:
176                ref_path = _GetResultAsXml(
177                    ref_path, os.path.join(temp_dir, "ref"))
178                if not ref_path:
179                    return False
180                with open(ref_path, "r") as ref_file:
181                    if args.primary_abi_only:
182                        ref_build_attrs = xml_utils.GetAttributes(
183                            ref_file, common._BUILD_TAG,
184                            (common._ABIS_ATTR_KEY, ))
185                        ref_file.seek(0)
186                        abis = ref_build_attrs[
187                            common._ABIS_ATTR_KEY].split(",")
188                        ref_bitness = str(result_utils.GetAbiBitness(abis[0]))
189                        ref_dict = _LoadRefResults(ref_file, src_dict,
190                                                   ref_bitness, src_bitness)
191                    else:
192                        ref_dict = _LoadRefResults(ref_file, src_dict)
193        finally:
194            shutil.rmtree(temp_dir)
195
196        # Output
197        csv_file = string_io_module.StringIO()
198        try:
199            writer = csv.writer(csv_file, lineterminator="\n")
200
201            writer.writerows(extra_rows)
202
203            for keys, attrs in (
204                    (_RESULT_ATTR_KEYS, result_attrs),
205                    (_BUILD_ATTR_KEYS, build_attrs),
206                    (_SUMMARY_ATTR_KEYS, summary_attrs)):
207                writer.writerows((k, attrs.get(k, "")) for k in keys)
208
209            src_list = sorted(src_dict.items())
210            if ref_path:
211                _WriteComparisonToCsv(src_list, ref_dict, writer)
212            else:
213                _WriteResultsToCsv(src_list, writer)
214
215            if exceed_max:
216                writer.writerow((_TOO_MANY_DATA,))
217
218            client.import_csv(args.dest, csv_file.getvalue())
219        finally:
220            csv_file.close()
221
222
223def _DownloadResultZipFromGcs(gcs_url, local_dir):
224    """Downloads a result ZIP from GCS.
225
226    If the GCS URL is a directory, this function searches the directory for the
227    file whose name matches the pattern of result ZIP.
228
229    Args:
230        gcs_url: The URL of the ZIP file or the directory.
231        local_dir: The local directory where the ZIP is downloaded.
232
233    Returns:
234        The path to the downloaded ZIP file.
235        None if fail to download.
236    """
237    gsutil_path = gcs_utils.GetGsutilPath()
238    if not gsutil_path:
239        return False
240
241    if gcs_utils.IsGcsFile(gsutil_path, gcs_url):
242        gcs_urls = [gcs_url]
243    else:
244        ls_urls = gcs_utils.List(gsutil_path, gcs_url)
245        gcs_urls = [x for x in ls_urls if
246                    re.match(".+/results_\\d*\\.zip$", x)]
247        if not gcs_urls:
248            gcs_urls = [x for x in ls_urls if
249                        re.match(".+/log-result_\\d*\\.zip$", x)]
250
251    if not gcs_urls:
252        logging.error("No results on %s", gcs_url)
253        return None
254    if len(gcs_urls) > 1:
255        logging.warning("More than one result. Select %s", gcs_urls[0])
256
257    if not os.path.exists(local_dir):
258        os.makedirs(local_dir)
259    if not gcs_utils.Copy(gsutil_path, gcs_urls[0], local_dir):
260        logging.error("Fail to copy from %s", gcs_urls[0])
261        return None
262
263    return os.path.join(local_dir, gcs_urls[0].rpartition("/")[2])
264
265
266def _GetResultAsXml(src, temp_dir):
267    """Downloads and extracts an XML result.
268
269    If src is a GCS URL, it is downloaded to temp_dir.
270    If the file is a ZIP, it is extracted to temp_dir.
271
272    Args:
273        src: The location of the file, can be a directory on GCS,
274             a ZIP file on GCS, a local ZIP file, or a local XML file.
275        temp_dir: The directory where the ZIP is downloaded and extracted.
276
277    Returns:
278        The path to the XML file.
279        None if fails to the find the XML.
280    """
281    original_src = src
282    if src.startswith("gs://"):
283        src = _DownloadResultZipFromGcs(src, os.path.join(temp_dir, "zipped"))
284        if not src:
285            return None
286
287    if zipfile.is_zipfile(src):
288        with zipfile.ZipFile(src, mode="r") as zip_file:
289            src = result_utils.ExtractResultZip(
290                zip_file, os.path.join(temp_dir, "unzipped"))
291        if not src:
292            logging.error("Cannot find XML result in %s", original_src)
293            return None
294
295    return src
296
297
298def _FilterTestResults(xml_file, max_return, filter_func):
299    """Loads test results from XML to dictionary with a filter.
300
301    Args:
302        xml_file: The input file object in XML format.
303        max_return: Maximum number of output results.
304        filter_func: A function taking the test name and result as parameters,
305                     and returning whether it should be included.
306
307    Returns:
308        A dict of {name: result} where name is a tuple of strings and result
309        is a string.
310    """
311    result_dict = dict()
312    for module, testcase, test in result_utils.IterateTestResults(xml_file):
313        if len(result_dict) >= max_return:
314            break
315        test_name = result_utils.GetTestName(module, testcase, test)
316        result = test.attrib.get(common._RESULT_ATTR_KEY, "")
317        if filter_func(test_name, result):
318            result_dict[test_name] = result
319
320    return result_dict
321
322
323def _LoadSrcResults(src_xml, max_return, bitness=""):
324    """Loads test results from XML to dictionary.
325
326    If number of results exceeds max_return, only failures are returned.
327    If number of failures exceeds max_return, the results are truncated.
328
329    Args
330        src_xml: The file object in XML format.
331        max_return: Maximum number of returned results.
332        bitness: A string, the bitness of the returned results.
333
334    Returns:
335        A dict of {name: result} and a boolean which represents whether the
336        results are truncated.
337    """
338    def FilterBitness(name):
339        return not bitness or bitness == name[0]
340
341    results = _FilterTestResults(
342        src_xml, max_return + 1, lambda name, result: FilterBitness(name))
343
344    if len(results) > max_return:
345        src_xml.seek(0)
346        results = _FilterTestResults(
347            src_xml, max_return + 1,
348            lambda name, result: result == "fail" and FilterBitness(name))
349
350    exceed_max = len(results) > max_return
351    if results and exceed_max:
352        del results[max(results)]
353
354    return results, exceed_max
355
356
357def _LoadRefResults(ref_xml, base_results, ref_bitness="", base_bitness=""):
358    """Loads reference results from XML to dictionary.
359
360    A test result in ref_xml is returned if the test fails in base_results.
361
362    Args:
363        ref_xml: The file object in XML format.
364        base_results: A dict of {name: result} containing the test names to be
365                      loaded from ref_xml.
366        ref_bitness: A string, the bitness of the results to be loaded from
367                     ref_xml.
368        base_bitness: A string, the bitness of the returned results. If this
369                      argument is specified, the function ignores bitness when
370                      comparing test names.
371
372    Returns:
373        A dict of {name: result}, the test name in base_results and the result
374        in ref_xml.
375    """
376    ref_results = dict()
377    for module, testcase, test in result_utils.IterateTestResults(ref_xml):
378        if len(ref_results) >= len(base_results):
379            break
380        result = test.attrib.get(common._RESULT_ATTR_KEY, "")
381        name = result_utils.GetTestName(module, testcase, test)
382
383        if ref_bitness and name[0] != ref_bitness:
384            continue
385        if base_bitness:
386            name_in_base = (base_bitness, ) + name[1:]
387        else:
388            name_in_base = name
389
390        if base_results.get(name_in_base, "") == "fail":
391            ref_results[name_in_base] = result
392
393    return ref_results
394
395
396def _WriteResultsToCsv(result_list, writer):
397    """Writes a list of test names and results to a CSV file.
398
399    Args:
400        result_list: The list of (name, result).
401        writer: The object of CSV writer.
402    """
403    writer.writerow(_TABLE_HEADER)
404    writer.writerows(name + (result,) for name, result in result_list)
405
406
407def _WriteComparisonToCsv(result_list, reference_dict, writer):
408    """Writes test names, results, and reference results to a CSV file.
409
410    Args:
411        result_list: The list of (name, result).
412        reference_dict: The dict of {name: reference_result}.
413        writer: The object of CSV writer.
414    """
415    writer.writerow(_CMP_TABLE_HEADER)
416    for name, result in result_list:
417        if result == "fail":
418            reference = reference_dict.get(name, "no_data")
419        else:
420            reference = ""
421        writer.writerow(name + (result, reference))
422