tools/wrap-ci-test.py

   1 #!/usr/bin/env python3
   2
   3 #
   4 # Add arbritrary commands to a GitLab CI compatible (JUnit) test report
   5 # SPDX-License-Identifier: MIT
   6 #
   7 # Usage:
   8 #   wrap-ci-test --file foo.xml --suite "Suite" --case "Name" --command "command"
   9 #   wrap-ci-test --file foo.xml --suite "Suite" --case "Name" command [args] ...
  10
  11 # This script runs a command and adds it to a JUnit report which can then
  12 # be used as a GitLab CI test report:
  13 #
  14 #   https://docs.gitlab.com/ee/ci/testing/unit_test_reports.html
  15 #
  16 # Commands can be specified with the "--command" flag, which will run
  17 # in a subshell, or as a list of extra arguments, which will be run
  18 # directly.
  19 #
  20 # Command output will be "teed". Scrubbed versions will be added to the
  21 # report and unmodified versions will be printed to stdout and stderr.
  22 #
  23 # If the command exit code is nonzero it will be added to the report
  24 # as a failure.
  25 #
  26 # The wrapper will return the command exit code.
  27
  28 # JUnit report information can be found at
  29 # https://github.com/testmoapp/junitxml
  30 # https://www.ibm.com/docs/en/developer-for-zos/14.2?topic=formats-junit-xml-format
  31
  32
  33 import argparse
  34 import html
  35 import time
  36 import pathlib
  37 import re
  38 import subprocess
  39 import sys
  40 import xml.etree.ElementTree as ET
  41
  42
  43 def main():
  44     parser = argparse.ArgumentParser(usage='\n  %(prog)s [options] --command "command"\n  %(prog)s [options] command ...')
  45     parser.add_argument('--file', required=True, type=pathlib.Path, help='The JUnit-compatible XML file')
  46     parser.add_argument('--suite', required=True, help='The testsuite_el name')
  47     parser.add_argument('--case', required=True, help='The testcase name')
  48     parser.add_argument('--command', help='The command to run if no extra arguments are provided')
  49
  50     args, command_list = parser.parse_known_args()
  51
  52     if (args.command and len(command_list) > 0) or (args.command is None and len(command_list) == 0):
  53         sys.stderr.write('Error: The command must be provided via the --command flag or extra arguments.\n')
  54         sys.exit(1)
  55
  56     try:
  57         tree = ET.parse(args.file)
  58         testsuites_el = tree.getroot()
  59     except FileNotFoundError:
  60         testsuites_el = ET.Element('testsuites')
  61         tree = ET.ElementTree(testsuites_el)
  62     except ET.ParseError:
  63         sys.stderr.write(f'Error: {args.file} is invalid.\n')
  64         sys.exit(1)
  65
  66     suites_time = float(testsuites_el.get('time', 0.0))
  67     suites_tests = int(testsuites_el.get('tests', 0)) + 1
  68     suites_failures = int(testsuites_el.get('failures', 0))
  69
  70     testsuite_el = testsuites_el.find(f'./testsuite[@name="{args.suite}"]')
  71     if testsuite_el is None:
  72         testsuite_el = ET.Element('testsuite', attrib={'name': args.suite})
  73         testsuites_el.append(testsuite_el)
  74
  75     suite_time = float(testsuite_el.get('time', 0.0))
  76     suite_tests = int(testsuite_el.get('tests', 0)) + 1
  77     suite_failures = int(testsuite_el.get('failures', 0))
  78
  79     testcase_el = ET.Element('testcase', attrib={'name': args.case})
  80     testsuite_el.append(testcase_el)
  81
  82     if args.command:
  83         proc_args = args.command
  84         in_shell = True
  85     else:
  86         proc_args = command_list
  87         in_shell = False
  88
  89     start_time = time.perf_counter()
  90     proc = subprocess.run(proc_args, shell=in_shell, encoding='UTF-8', errors='replace', capture_output=True)
  91     case_time = time.perf_counter() - start_time
  92
  93     testcase_el.set('time', f'{case_time}')
  94     testsuite_el.set('time', f'{suite_time + case_time}')
  95     testsuites_el.set('time', f'{suites_time + case_time}')
  96
  97     # XXX Try to interleave them?
  98     sys.stdout.write(proc.stdout)
  99     sys.stderr.write(proc.stderr)
 100
 101     # Remove ANSI control sequences and escape other invalid characters
 102     # https://stackoverflow.com/a/14693789/82195
 103     ansi_seq_re = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
 104     scrubbed_stdout = html.escape(ansi_seq_re.sub('', proc.stdout), quote=False)
 105     scrubbed_stderr = html.escape(ansi_seq_re.sub('', proc.stderr), quote=False)
 106
 107     if proc.returncode != 0:
 108         failure_el = ET.Element('failure')
 109         failure_el.text = f'{scrubbed_stdout}{scrubbed_stderr}'
 110         testcase_el.append(failure_el)
 111         testsuite_el.set('failures', f'{suite_failures + 1}')
 112         testsuites_el.set('failures', f'{suites_failures + 1}')
 113     else:
 114         system_out_el = ET.Element('system-out')
 115         system_out_el.text = f'{scrubbed_stdout}'
 116         testcase_el.append(system_out_el)
 117         system_err_el = ET.Element('system-err')
 118         system_err_el.text = f'{scrubbed_stderr}'
 119         testcase_el.append(system_err_el)
 120
 121     testsuite_el.set('tests', f'{suite_tests}')
 122     testsuites_el.set('tests', f'{suites_tests}')
 123
 124     tree.write(args.file, encoding='UTF-8', xml_declaration=True)
 125
 126     return proc.returncode
 127
 128 if __name__ == '__main__':
 129     sys.exit(main())