mongo/buildscripts/antithesis_suite.py

#!/usr/bin/env python3
"""Command line utility for generating suites for targeting antithesis."""

import os.path
import sys

import click
import yaml

# Get relative imports to work when the package is not installed on the PYTHONPATH.
if __name__ == "__main__" and __package__ is None:
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

HOOKS_BLACKLIST = [
    "CleanEveryN",
    "ContinuousStepdown",
    "CheckOrphansDeleted",
    # TODO SERVER-70396 re-enable hook once the checkMetadata feature flag is removed
    # To check the feature flag we need to contact directly the config server that is not exposed in the ExternalFixture
    "CheckMetadataConsistencyInBackground",
]

_SUITES_PATH = os.path.join("buildscripts", "resmokeconfig", "suites")

MONGOS_PORT = 27017


def delete_archival(suite):
    """Remove archival for Antithesis environment."""
    suite.pop("archive", None)
    suite.get("executor", {}).pop("archive", None)


def make_hooks_compatible(suite):
    """Make hooks compatible in Antithesis environment."""
    if suite.get("executor", {}).get("hooks", None):
        # it's either a list of strings, or a list of dicts, each with key 'class'
        if isinstance(suite["executor"]["hooks"][0], str):
            suite["executor"]["hooks"] = ["AntithesisLogging"] + [
                hook for hook in suite["executor"]["hooks"] if hook not in HOOKS_BLACKLIST
            ]
        elif isinstance(suite["executor"]["hooks"][0], dict):
            suite["executor"]["hooks"] = [{"class": "AntithesisLogging"}] + [
                hook for hook in suite["executor"]["hooks"] if hook["class"] not in HOOKS_BLACKLIST
            ]
        else:
            raise RuntimeError('Unknown structure in hook. File a TIG ticket.')


def use_external_fixture(suite_name, suite):
    """Use external version of this fixture."""
    if suite.get("executor", {}).get("fixture", None):
        suite["executor"]["fixture"] = {
            "class": f"External{suite['executor']['fixture']['class']}",
            "original_suite_name": suite_name,
        }


def get_mongos_connection_url(suite):
    """
    Return the mongos connection URL for suite if Antithesis compatible.

    :param suite: Parsed YAML document for the suite we wish to connect to.
    :return: Connection url for the suite, or a warning if Antithesis incompatible.
    """
    if suite.get("executor", {}).get("fixture", {}).get("num_mongos", None):
        return "mongodb://" + ",".join(
            [f"mongos{i}:{MONGOS_PORT}" for i in range(suite['executor']['fixture']['num_mongos'])])
    else:
        return "ANTITHESIS_INCOMPATIBLE"


def update_test_data(suite):
    """Update TestData to be compatible with antithesis."""
    suite.setdefault("executor", {}).setdefault(
        "config", {}).setdefault("shell_options", {}).setdefault("global_vars", {}).setdefault(
            "TestData", {}).update({"useActionPermittedFile": False})


def update_shell(suite):
    """Update shell for when running in Antithesis."""
    suite.setdefault("executor", {}).setdefault("config", {}).setdefault("shell_options",
                                                                         {}).setdefault("eval", "")
    suite["executor"]["config"]["shell_options"]["eval"] += "jsTestLog = Function.prototype;"


def update_exclude_tags(suite):
    """Update the exclude tags to exclude antithesis incompatible tests."""
    suite.setdefault('selector', {})
    if not suite.get('selector').get('exclude_with_any_tags'):
        suite['selector']['exclude_with_any_tags'] = ["antithesis_incompatible"]
    else:
        suite['selector']['exclude_with_any_tags'].append('antithesis_incompatible')


def get_antithesis_suite_config(suite_name):
    """Modify suite in-place to be antithesis compatible."""
    with open(os.path.join(_SUITES_PATH, f"{suite_name}.yml")) as fstream:
        suite = yaml.safe_load(fstream)

    delete_archival(suite)
    make_hooks_compatible(suite)
    use_external_fixture(suite_name, suite)
    update_test_data(suite)
    update_shell(suite)
    update_exclude_tags(suite)

    return suite


@click.group()
def cli():
    """CLI Entry point."""
    pass


def _generate(suite_name: str) -> None:
    suite = get_antithesis_suite_config(suite_name)

    out = yaml.dump(suite)
    with open(os.path.join(_SUITES_PATH, f"antithesis_{suite_name}.yml"), "w") as fstream:
        fstream.write(
            "# this file was generated by buildscripts/antithesis_suite.py generate {}\n".format(
                suite_name))
        fstream.write("# Do not modify by hand\n")
        fstream.write(out)


@cli.command()
@click.argument('suite_name')
def generate(suite_name: str) -> None:
    """Generate a single suite."""
    _generate(suite_name)


@cli.command('generate-all')
def generate_all():
    """Generate all suites."""
    for path in os.listdir(_SUITES_PATH):
        if os.path.isfile(os.path.join(_SUITES_PATH, path)):
            suite = path.split(".")[0]
            _generate(suite)


if __name__ == "__main__":
    cli()