testing_support/fake_cipd.py

#!/usr/bin/env python3
# Copyright (c) 2018 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
from string import Template

import argparse
import io
import json
import os
import re
import shutil
import sys

ARCH_VAR = 'arch'
OS_VAR = 'os'
PLATFORM_VAR = 'platform'

CIPD_SUBDIR_RE = '@Subdir (.*)'
CIPD_DESCRIBE = 'describe'
CIPD_ENSURE = 'ensure'
CIPD_ENSURE_FILE_RESOLVE = 'ensure-file-resolve'
CIPD_EXPAND_PKG = 'expand-package-name'
CIPD_EXPORT = 'export'

DESCRIBE_STDOUT_TEMPLATE = """\
Package:       ${package}
Instance ID:   ${package}-fake-instance-id
Registered by: user:fake-testing-support
Registered at: 2023-05-10 18:53:55.078574 +0000 UTC
Refs:
  ${package}-latest
Tags:
  git_revision:${package}-fake-git-revision
  ${package}-fake-tag:1.0
  ${package}-fake-tag:2.0
"""

DESCRIBE_JSON_TEMPLATE = """{
  "result": {
    "pin": {
      "package": "${package}",
      "instance_id": "${package}-fake-instance-id"
    },
    "registered_by": "user:fake-testing-support",
    "registered_ts": 1683744835,
    "refs": [
      {
        "ref": "${package}-latest",
        "instance_id": "${package}-fake-instance-id",
        "modified_by": "user:fake-testing-support",
        "modified_ts": 1683744835
      }
    ],
    "tags": [
      {
        "tag": "git_revision:${package}-fake-git-revision",
        "registered_by": "user:fake-testing-support",
        "registered_ts": 1683744835
      },
      {
        "tag": "${package}-fake-tag:1.0",
        "registered_by": "user:fake-testing-support",
        "registered_ts": 1683744835
      },
      {
        "tag": "${package}-fake-tag:2.0",
        "registered_by": "user:fake-testing-support",
        "registered_ts": 1683744835
      }
    ]
  }
}"""


def parse_cipd(root, contents):
    tree = {}
    current_subdir = None
    for line in contents:
        line = line.strip()
        match = re.match(CIPD_SUBDIR_RE, line)
        if match:
            print('match')
            current_subdir = os.path.join(root, *match.group(1).split('/'))
            if not root:
                current_subdir = match.group(1)
        elif line and current_subdir:
            print('no match')
            tree.setdefault(current_subdir, []).append(line)
    return tree


def expand_package_name_cmd(package_name):
    package_split = package_name.split("/")
    suffix = package_split[-1]
    # Any use of var equality should return empty for testing.
    if "=" in suffix:
        if suffix != "${platform=fake-platform-ok}":
            return ""
        package_name = "/".join(package_split[:-1] + ["${platform}"])
    for v in [ARCH_VAR, OS_VAR, PLATFORM_VAR]:
        var = "${%s}" % v
        if package_name.endswith(var):
            package_name = package_name.replace(var,
                                                "%s-expanded-test-only" % v)
    return package_name


def ensure_file_resolve():
    resolved = {"result": {}}
    parser = argparse.ArgumentParser()
    parser.add_argument('-ensure-file', required=True)
    parser.add_argument('-json-output')
    args, _ = parser.parse_known_args()
    with io.open(args.ensure_file, 'r', encoding='utf-8') as f:
        new_content = parse_cipd("", f.readlines())
        for path, packages in new_content.items():
            resolved_packages = []
            for package in packages:
                package_name = expand_package_name_cmd(package.split(" ")[0])
                resolved_packages.append({
                    "package": package_name,
                    "pin": {
                        "package": package_name,
                        "instance_id": package_name + "-fake-resolved-id",
                    }
                })
            resolved["result"][path] = resolved_packages
        with io.open(args.json_output, 'w', encoding='utf-8') as f:
            f.write(json.dumps(resolved, indent=4))


def describe_cmd(package_name):
    parser = argparse.ArgumentParser()
    parser.add_argument('-json-output')
    parser.add_argument('-version', required=True)
    args, _ = parser.parse_known_args()
    json_template = Template(DESCRIBE_JSON_TEMPLATE).substitute(
        package=package_name)
    cli_out = Template(DESCRIBE_STDOUT_TEMPLATE).substitute(
        package=package_name)
    json_out = json.loads(json_template)
    found = False
    for tag in json_out['result']['tags']:
        if tag['tag'] == args.version:
            found = True
            break
    for tag in json_out['result']['refs']:
        if tag['ref'] == args.version:
            found = True
            break
    if found:
        if args.json_output:
            with io.open(args.json_output, 'w', encoding='utf-8') as f:
                f.write(json.dumps(json_out, indent=4))
        sys.stdout.write(cli_out)
        return 0
    sys.stdout.write('Error: no such ref.\n')
    return 1


def main():
    cmd = sys.argv[1]
    assert cmd in [
        CIPD_DESCRIBE, CIPD_ENSURE, CIPD_ENSURE_FILE_RESOLVE, CIPD_EXPAND_PKG,
        CIPD_EXPORT
    ]
    # Handle cipd expand-package-name
    if cmd == CIPD_EXPAND_PKG:
        # Expecting argument after cmd
        assert len(sys.argv) == 3
        # Write result to stdout
        sys.stdout.write(expand_package_name_cmd(sys.argv[2]))
        return 0
    if cmd == CIPD_DESCRIBE:
        # Expecting argument after cmd
        assert len(sys.argv) >= 3
        return describe_cmd(sys.argv[2])
    if cmd == CIPD_ENSURE_FILE_RESOLVE:
        return ensure_file_resolve()

    parser = argparse.ArgumentParser()
    parser.add_argument('-ensure-file')
    parser.add_argument('-root')
    args, _ = parser.parse_known_args()

    with io.open(args.ensure_file, 'r', encoding='utf-8') as f:
        new_content = parse_cipd(args.root, f.readlines())

    # Install new packages
    for path, packages in new_content.items():
        if not os.path.exists(path):
            os.makedirs(path)
        with io.open(os.path.join(path, '_cipd'), 'w', encoding='utf-8') as f:
            f.write('\n'.join(packages))

    # Save the ensure file that we got
    shutil.copy(args.ensure_file, os.path.join(args.root, '_cipd'))

    return 0


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