squad-download-attachments

#!/usr/bin/env python3
# vim: set ts=4
#
# Copyright 2023-present Linaro Limited
#
# SPDX-License-Identifier: MIT

import argparse
import contextlib
import csv
import glob
import json
import logging
import lzma
import re
import statistics
import sys
import tarfile
from os import chdir
from pathlib import Path

from squad_client.core.api import SquadApi
from squad_client.core.models import ALL, Environment, Squad, TestRun
from squad_client.shortcuts import download_attachments
from squad_client.utils import getid

SquadApi.configure(cache=3600, url="https://qa-reports.linaro.org/")

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def arg_parser():
    parser = argparse.ArgumentParser(description="download attachments in a testrun")

    parser.add_argument(
        "--group",
        required=False,
        help="squad group",
    )

    parser.add_argument(
        "--project",
        required=False,
        help="squad project",
    )

    parser.add_argument(
        "--build",
        required=False,
        help="squad build",
    )

    parser.add_argument(
        "--csv",
        default=False,
        action="store_true",
        required=False,
        help="generate csv files",
    )

    parser.add_argument(
        "--path",
        required=False,
        help="tuxrun artefact path",
    )

    return parser


def generate_files(dirname, fileprefix, create_csv):
    file = glob.glob(f"{dirname}/mmtests-*.tar.xz")
    # Extract the json file that contains the benchmark data.
    with contextlib.closing(lzma.LZMAFile(file[0])) as xz:
        with tarfile.open(fileobj=xz) as f:
            f.extractall(dirname)

    file = glob.glob(f"{dirname}/output/*.json")
    filename = Path(file[0]).name
    filename = re.sub(r"^.*CONFIG", fileprefix, filename)
    file_write = f"{dirname}/{filename}"
    # sort the json keys in the benchmark data file.
    with open(file[0], mode="r") as read_file:
        pjson = json.dumps(json.load(read_file), sort_keys=True, indent=4)
        with open(file_write, mode="w") as write_file:
            write_file.write(pjson)

        if not create_csv:
            return True

        with open(file_write.replace(".json", ".csv"), mode="w") as csv_file:
            csv_writer = csv.writer(csv_file)
            dict_json = json.loads(pjson)
            if not dict_json["results"]:
                return False
            headers = [
                "median",
                "average",
                "standard deviation",
                "Percentage of std",
                "name",
                "iteration",
                "name_iteration",
                "raw data...",
            ]
            csv_writer.writerow(headers)
            for key in dict_json["results"]["_ResultData"]:
                iterations = 0
                for k in dict_json["results"]["_ResultData"][key]:
                    csv_data = []
                    float_arr = []
                    for number in k["Values"]:
                        float_arr.append(float(number))
                    csv_data.append(statistics.median(float_arr))
                    csv_data.append(statistics.mean(float_arr))
                    csv_data.append(statistics.stdev(float_arr))
                    if statistics.mean(float_arr) == 0:
                        continue
                    csv_data.append(
                        (statistics.stdev(float_arr) / statistics.mean(float_arr)) * 100
                    )
                    csv_data.append(key)
                    iterations = iterations + 1
                    csv_data.append(f"iteration_{iterations}")
                    csv_data.append(f"{key}_iteration_{iterations}")
                    csv_data.extend(k["Values"])
                    csv_writer.writerow(csv_data)
    return True


def run():
    args = arg_parser().parse_args()

    if args.path:
        generate_files(args.path, "local-", args.csv)
    else:
        group = Squad().group(args.group)
        if group is None:
            logger.error(f"Get group failed. Group not found: '{args.group}'.")
            return -1

        project = group.project(args.project)
        if project is None:
            logger.error(f"Get project failed. Project not found: '{args.project}'.")
            return -1

        build = project.build(args.build)
        if build is None:
            logger.error(f"Get build failed. Build not found: '{args.build}'.")
            return -1

        environments = project.environments(count=ALL, ordering="slug").values()
        if not environments:
            logger.error("Get environments failed. No environments found.")
            return -1

        suites = project.suites(count=ALL, ordering="slug").values()
        if not suites:
            logger.error("Get suites failed. No suites found.")
            return -1

        attachment_dir = Path("stored_attachments/" + args.build)
        testruns = build.testruns()
        for testrun in testruns:
            if not TestRun(testrun).attachments:
                continue
            env_name = Environment(getid((TestRun(testrun).environment))).slug
            dirname = Path(f"{attachment_dir}/{env_name}_{str(TestRun(testrun).id)}")
            print(dirname)
            # Only picking up 'qemu-' environments
            # The check will be 'not "build" in dirname.name' when DUT in tuxbridge supports attachments.
            if "qemu-" in dirname.name:
                Path.mkdir(dirname, parents=True, exist_ok=True)
                chdir(dirname)
                download_attachments(TestRun(testrun))
                chdir(sys.path[0])
                fileprefix = (
                    f"tux-{re.sub(r'_[0-9]+$', '-', dirname.name.replace('qemu-', ''))}"
                )

                generate_files(dirname, fileprefix, args.csv)


if __name__ == "__main__":
    sys.exit(run())