src/tests/restricted_traces/gen_restricted_traces.py

#!/usr/bin/python3
#
# Copyright 2020 The ANGLE Project Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
#
# gen_restricted_traces.py:
#   Generates integration code for the restricted trace tests.

import getpass
import glob
import fnmatch
import re
import json
import os
import sys
import subprocess

CIPD_TRACE_PREFIX = 'angle/traces'
EXPERIMENTAL_CIPD_PREFIX = 'experimental/google.com/%s/angle/traces'
DEPS_PATH = '../../../DEPS'
DEPS_START = '# === ANGLE Restricted Trace Generated Code Start ==='
DEPS_END = '# === ANGLE Restricted Trace Generated Code End ==='
DEPS_TEMPLATE = """\
  'src/tests/restricted_traces/{trace}': {{
      'packages': [
        {{
            'package': '{trace_prefix}/{trace}',
            'version': 'version:{version}',
        }},
      ],
      'dep_type': 'cipd',
      'condition': 'checkout_angle_restricted_trace_{trace}',
  }},
"""
STAR_PATH = '../../../infra/config/representative_traces.star'


DEPS_VAR_START = '# === ANGLE Restricted Trace Generated Var Start ==='
DEPS_VAR_END = '# === ANGLE Restricted Trace Generated Var End ==='
DEPS_VAR_TEMPLATE = """\
  'checkout_angle_restricted_trace_{trace}': '{condition}',
"""

def reject_duplicate_keys(pairs):
    found_keys = {}
    for key, value in pairs:
        if key in found_keys:
            raise ValueError("duplicate key: %r" % (key,))
        else:
            found_keys[key] = value
    return found_keys


def read_json(json_file):
    with open(json_file) as map_file:
        return json.loads(map_file.read(), object_pairs_hook=reject_duplicate_keys)


def write_json(json_file, data):
    with open(json_file, 'w') as map_file:
        json.dump(data, map_file, indent=4, sort_keys=True)


def update_deps(trace_pairs, representative_traces, cq_extra_traces):
    # Generate substitution string
    replacement = ""
    for (trace, version) in trace_pairs:
        if 'x' in version:
            version = version.strip('x')
            trace_prefix = EXPERIMENTAL_CIPD_PREFIX % getpass.getuser()
        else:
            trace_prefix = CIPD_TRACE_PREFIX
        sub = {'trace': trace, 'version': version, 'trace_prefix': trace_prefix}
        replacement += DEPS_TEMPLATE.format(**sub)

    # Update DEPS to download CIPD dependencies
    with open(DEPS_PATH) as f:
        lines = f.readlines()

    def slice_to_replace(lines, start_tag, end_tag):
        start, end = [i for i, s in enumerate(lines) if start_tag in s or end_tag in s]
        return slice(start + 1, end)

    # Replace lines between DEPS_START and DEPS_END with new code
    lines[slice_to_replace(lines, DEPS_START, DEPS_END)] = [replacement]

    # Replace lines between DEPS_VAR_START and DEPS_VAR_END with new code
    deps_vars = []
    for (trace, _) in trace_pairs:
        if trace in representative_traces or trace in cq_extra_traces:
            condition = 'checkout_angle_restricted_traces'
        else:
            condition = 'checkout_extra_traces'
        deps_vars.append(DEPS_VAR_TEMPLATE.format(trace=trace, condition=condition))

    lines[slice_to_replace(lines, DEPS_VAR_START, DEPS_VAR_END)] = deps_vars

    with open(DEPS_PATH, 'w') as f:
        f.write(''.join(lines))

    return True


def update_star(smoke_traces, representative_traces, cq_extra_traces):
    star_content = '''\
# Copyright 2026 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Generated by gen_restricted_traces.py. DO NOT EDIT."""

SMOKE_TRACES = [
{smoke_traces}]

REPRESENTATIVE_TRACES = [
{representative_traces}]
'''
    combined_representative = sorted(list(set(representative_traces + cq_extra_traces)))
    smoke_list = "".join([f'    "{trace}",\n' for trace in sorted(smoke_traces)])
    rep_list = "".join([f'    "{trace}",\n' for trace in combined_representative])

    with open(STAR_PATH, 'w') as f:
        f.write(star_content.format(smoke_traces=smoke_list, representative_traces=rep_list))

    return True


def run_lucicfg():
    import platform
    infra_config_dir = os.path.abspath('../../../infra/config')
    cmd = ['lucicfg']
    if platform.system() == 'Windows':
        cmd = ['lucicfg.bat']
    try:
        subprocess.check_output(
            cmd + ['generate', 'main.star'], cwd=infra_config_dir, stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError as e:
        print(f'lucicfg failed: {e.output.decode()}')
        return False
    except Exception as e:
        print(f'Error running lucicfg: {e}')
        return False
    return True


def main():
    json_file = 'restricted_traces.json'

    # Need to change directory to the script's directory so relative paths work
    # if run from other directories.
    os.chdir(os.path.dirname(os.path.abspath(__file__)))

    json_data = read_json(json_file)
    if 'traces' not in json_data:
        print('Trace data missing traces key.')
        return 1
    raw_trace_parts = [trace.split(' ') for trace in json_data['traces']]
    representative_traces = [p[0] for p in raw_trace_parts if 'representative' in p[2:]]
    smoke_traces = [p[0] for p in raw_trace_parts if 'smoke' in p[2:]]
    cq_extra_traces = [
        p[0] for p in raw_trace_parts if 'ci' not in p[2:] and 'representative' not in p[2:]
    ]
    trace_pairs = [(p[0], p[1]) for p in raw_trace_parts]


    # auto_script parameters.
    if len(sys.argv) > 1:
        inputs = [json_file]
        outputs = ['.gitignore', STAR_PATH]

        if sys.argv[1] == 'inputs':
            print(','.join(inputs))
        elif sys.argv[1] == 'outputs':
            print(','.join(outputs))
        else:
            print('Invalid script parameters.')
            return 1
        return 0

    LIMIT_N = 10
    if len(cq_extra_traces) > LIMIT_N:
        print(f'Error: Too many CQ extra traces ({len(cq_extra_traces)}). Limit is {LIMIT_N}.')
        print('Please move some traces to conditional checkout by adding the "ci" tag.')
        return 1

    if not update_deps(trace_pairs, representative_traces, cq_extra_traces):
        print('DEPS file update failed')
        return 1

    if not update_star(smoke_traces, representative_traces, cq_extra_traces):
        print('representative_traces.star update failed')
        return 1

    if not run_lucicfg():
        print('lucicfg generate failed')
        return 1

    return 0


if __name__ == '__main__':
    sys.exit(main())