tools/assistant/instrument_with_assistant_and_run_test.py

#!/usr/bin/env python3
# Copyright (c) 2025, the R8 project authors. Please see the AUTHORS file
# for details. All rights reserved. Use of this source code is governed by a
# BSD-style license that can be found in the LICENSE file.

import argparse
import os
import shutil
import sys

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import adb
import apk_masseur
import compiledump
import utils


class AttrDict(dict):

    def __getattr__(self, name):
        return self.get(name, None)


def run_r8_assistant(options, temp_dir):
    compiledump_args = AttrDict(
        vars(
            argparse.Namespace(
                dump=options.dump,
                compiler='assistant',
                disable_assertions=options.disable_assertions,
                debug_agent=options.debug_agent,
                version='main',
                no_build=options.no_build,
                enable_missing_library_api_modeling=False,
                java_opts=[],
            )))

    instrumented_dex = os.path.join(temp_dir, 'out.jar')

    with utils.TempDir() as compiledump_temp_dir:
        compile_result = compiledump.run1(compiledump_temp_dir,
                                          compiledump_args,
                                          ['--output', instrumented_dex])
        if compile_result != 0:
            raise Exception('Failed to run R8 assistant')

    original_app_apk = options.apk

    # Create new app apk with instrumented dex.
    instrumented_app_apk = os.path.join(temp_dir, "app.apk")
    apk_masseur.masseur(original_app_apk,
                        dex=instrumented_dex,
                        out=instrumented_app_apk,
                        keystore=options.keystore)

    # Sign test apk with same key as the instrumented dex apk.
    original_test_apk = options.apk_test
    test_apk_destination = os.path.join(temp_dir, "test.apk")
    apk_masseur.masseur(original_test_apk,
                        out=test_apk_destination,
                        keystore=options.keystore)

    adb.run_instrumented(options.id,
                         options.id_test,
                         options.device,
                         instrumented_app_apk,
                         test_apk_destination,
                         quiet=False,
                         enable_logging=True,
                         test_runner=None)


def parse_options(argv):
    result = argparse.ArgumentParser(description='Run/compile dump artifacts.')
    result.add_argument('--apk', help='Path to apk for app')
    result.add_argument('--apk-test', help='Path to apk for androidTest')
    result.add_argument('--dump', help='Path to compiler dump')
    result.add_argument('--id', help='The id of the app')
    result.add_argument('--id-test', help='The id of the test')
    result.add_argument('--device', help='The device to run on')
    result.add_argument('--debug-agent',
                        help='Enable Java debug agent and suspend compilation ',
                        default=False,
                        action='store_true')
    result.add_argument(
        '--disable-assertions',
        '--disable_assertions',
        '-da',
        help='Disable Java assertions when running the compiler',
        default=False,
        action='store_true')
    result.add_argument('--keystore',
                        help='Path to app.keystore',
                        default=os.path.join(utils.TOOLS_DIR, 'debug.keystore'))
    result.add_argument('--no-build',
                        '--no_build',
                        help='Run without building first (only when using ToT)',
                        default=False,
                        action='store_true')
    result.add_argument('--temp',
                        help='A directory to use for temporaries and outputs.')
    return result.parse_known_args(argv)


def main(argv):
    options, _ = parse_options(argv)

    with utils.TempDir() as temp_dir:
        if options.temp:
            temp_dir = options.temp
            os.makedirs(temp_dir, exist_ok=True)
        run_r8_assistant(options, temp_dir)


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))