tests

#!/usr/bin/env python3
# Copyright (c) 2024 The mlkem-native project authors
# SPDX-License-Identifier: Apache-2.0

import platform
import os
import sys
import click
from functools import reduce

sys.path.append(f"{os.path.join(os.path.dirname(__file__), 'lib')}")
from mlkem_test import *
from util import path

"""
Click interface configuration
"""


def __callback(n):
    def callback(ctx, param, value):
        state = ctx.ensure_object(Options)
        state.__dict__[n] = value
        return value

    return callback


def add_options(options):
    """Shortcurt for adding multiple options for Click command"""
    return lambda func: reduce(lambda f, o: o(f), reversed(options), func)


_shared_options = [
    click.option(
        "-v",
        "--verbose",
        expose_value=False,
        is_flag=True,
        show_default=True,
        default=False,
        type=bool,
        help="Show verbose output or not",
        callback=__callback("verbose"),
    ),
    click.option(
        "-cp",
        "--cross-prefix",
        expose_value=False,
        default="",
        show_default=True,
        nargs=1,
        help="Cross prefix for compilation",
        callback=__callback("cross_prefix"),
    ),
    click.option(
        "--cflags",
        expose_value=False,
        nargs=1,
        help="Extra cflags to passed in (e.g. '-mcpu=cortex-a72')",
        callback=__callback("cflags"),
    ),
    click.option(
        "--arch-flags",
        expose_value=False,
        nargs=1,
        help="Extra arch flags to passed in (e.g. '-march=armv8')",
        callback=__callback("arch_flags"),
    ),
    click.option(
        "--auto/--no-auto",
        expose_value=False,
        is_flag=True,
        show_default=True,
        default=True,
        help="Allow makefile to auto configure system specific preprocessor",
        callback=__callback("auto"),
    ),
    click.option(
        "--opt",
        expose_value=False,
        nargs=1,
        type=click.Choice(["ALL", "OPT", "NO_OPT"], case_sensitive=False),
        show_default=True,
        default="ALL",
        help="Determine whether to compile/run the opt/no_opt binary or both",
        callback=__callback("opt"),
    ),
    click.option(
        "--compile/--no-compile",
        expose_value=False,
        is_flag=True,
        show_default=True,
        default=True,
        help="Determine to compile the binary or not",
        callback=__callback("compile"),
    ),
    click.option(
        "--run/--no-run",
        expose_value=False,
        is_flag=True,
        show_default=True,
        default=True,
        help="Determine to run the compiled binary or not",
        callback=__callback("run"),
    ),
    click.option(
        "-w",
        "--exec-wrapper",
        expose_value=False,
        show_default=True,
        default="",
        help="Run the benchmark binary with the user-customized wrapper.",
        callback=__callback("exec_wrapper"),
    ),
    click.option(
        "-r",
        "--run-as-root",
        expose_value=False,
        is_flag=True,
        show_default=True,
        default=False,
        type=bool,
        help="Benchmarking binary is run with sudo.",
        callback=__callback("run_as_root"),
    ),
]

_bench_options = [
    click.option(
        "-c",
        "--cycles",
        nargs=1,
        type=click.Choice(["NO", "PMU", "PERF", "M1"]),
        show_default=True,
        default="NO",
        help="Method for counting clock cycles. PMU requires (user-space) access to the Arm Performance Monitor Unit (PMU). PERF requires a kernel with perf support. M1 only works on Apple silicon.",
    ),
    click.option(
        "-o",
        "--output",
        nargs=1,
        help="Path to output file in json format",
    ),
    click.option(
        "-t",
        "--mac-taskpolicy",
        nargs=1,
        type=click.Choice(["utility", "background", "maintenance"]),
        hidden=platform.system() != "Darwin",
        show_default=True,
        default=None,
        help="Run the program using the specified QoS clamp. Applies to MacOS only. Setting this flag to 'background' guarantees running on E-cores. This is an abbreviation of --exec-wrapper 'taskpolicy -c {mac_taskpolicy}'.",
    ),
    click.option(
        "--components",
        is_flag=True,
        type=bool,
        show_default=True,
        default=False,
        help="Benchmark low-level components",
    ),
]


@click.group(invoke_without_command=True)
def cli():
    pass


@cli.command(
    short_help="Run the functional tests for all parameter sets",
    context_settings={"show_default": True},
)
@add_options(_shared_options)
@click.make_pass_decorator(Options, ensure=True)
def func(opts: Options):
    Tests(opts).func()


@cli.command(
    short_help="Run the nistkat tests for all parameter sets",
    context_settings={"show_default": True},
)
@add_options(_shared_options)
@click.make_pass_decorator(Options, ensure=True)
def nistkat(opts: Options):
    Tests(opts).nistkat()


@cli.command(
    short_help="Run the kat tests for all parameter sets",
    context_settings={"show_default": True},
)
@add_options(_shared_options)
@click.make_pass_decorator(Options, ensure=True)
def kat(opts: Options):
    Tests(opts).kat()


@cli.command(
    short_help="Run ACVP client",
    context_settings={"show_default": True},
)
@add_options(_shared_options)
@add_options(
    [
        click.option(
            "-d",
            "--acvp_dir",
            nargs=1,
            show_default=True,
            type=click.Path(),
            default=path("test/acvp_data"),
            help="Path to acvp directory",
        ),
    ]
)
@click.make_pass_decorator(Options, ensure=True)
def acvp(opts: Options, acvp_dir: str):
    Tests(opts).acvp(acvp_dir)


@cli.command(
    short_help="Run the benchmarks for all parameter sets",
    context_settings={"show_default": True},
)
@add_options(_shared_options)
@add_options(_bench_options)
@click.make_pass_decorator(Options, ensure=True)
def bench(
    opts: Options,
    cycles: str,
    output,
    mac_taskpolicy,
    components,
):
    Tests(opts).bench(
        cycles,
        output,
        mac_taskpolicy,
        components,
    )


@cli.command(
    short_help="Run all tests (except benchmark for now)",
    context_settings={"show_default": True},
)
@add_options(_shared_options)
@add_options(
    [
        click.option(
            "--func/--no-func",
            is_flag=True,
            show_default=True,
            default=True,
            help="Determine whether to run func test or not",
        ),
        click.option(
            "--kat/--no-kat",
            is_flag=True,
            show_default=True,
            default=True,
            help="Determine whether to run kat test or not",
        ),
        click.option(
            "--nistkat/--no-nistkat",
            is_flag=True,
            show_default=True,
            default=True,
            help="Determine whether to run nistkat test or not",
        ),
        click.option(
            "--acvp/--no-acvp",
            is_flag=True,
            show_default=True,
            default=True,
            help="Determine whether to run acvp test or not",
        ),
    ]
)
@click.make_pass_decorator(Options, ensure=True)
def all(
    opts: Options,
    func: bool,
    kat: bool,
    nistkat: bool,
    acvp: bool,
):
    Tests(opts).all(func, kat, nistkat, acvp)


@cli.command(
    short_help="Run the CBMC proofs for all parameter sets",
    context_settings={"show_default": True},
)
@click.make_pass_decorator(Options, ensure=True)
@add_options(
    [
        click.option(
            "--k",
            expose_value=False,
            nargs=1,
            type=click.Choice(["2", "3", "4", "ALL"]),
            show_default=True,
            default="ALL",
            help="MLKEM parameter set (MLKEM_K).",
            callback=__callback("k"),
        )
    ]
)
def cbmc(opts: Options):
    Tests(opts).cbmc(opts.k)


if __name__ == "__main__":
    cli()