robot-systest/autogeneration/generate_osm_test.py

#!/usr/bin/python3
# Copyright ETSI Contributors and Others.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.


# generate_osm_test.py --config test_config.yaml

import argparse
import logging
import os
import yaml

# from jinja2 import Environment, FileSystemLoader, select_autoescape
from jinja2 import Template


####################################
# Global functions
####################################
def set_logger(verbose):
    global logger
    log_format_simple = "%(levelname)s %(message)s"
    log_format_complete = "%(asctime)s %(levelname)s %(name)s %(filename)s:%(lineno)s %(funcName)s(): %(message)s"
    log_formatter_simple = logging.Formatter(
        log_format_simple, datefmt="%Y-%m-%dT%H:%M:%S"
    )
    handler = logging.StreamHandler()
    handler.setFormatter(log_formatter_simple)
    logger = logging.getLogger("generate_osm_test")
    logger.setLevel(level=logging.WARNING)
    logger.addHandler(handler)
    if verbose == 1:
        logger.setLevel(level=logging.INFO)
    elif verbose > 1:
        log_formatter = logging.Formatter(
            log_format_complete, datefmt="%Y-%m-%dT%H:%M:%S"
        )
        handler.setFormatter(log_formatter)
        logger.setLevel(level=logging.DEBUG)


def exit_with_error(message, code=1):
    logger.error(message)
    exit(code)


def render_template(template_file, test_data, output_file=None):
    logger.info(f"Rendering {template_file} with test data")

    # Load Jinja template
    with open(template_file, "r") as template_stream:
        template = Template(template_stream.read())

    # Render template with test_data and store it in output_file
    output = template.render(test_data)
    if output_file:
        if os.path.exists(output_file):
            exit_with_error(
                f"Output file '{output_file}' already exists. Use a different filename or remove it."
            )
        with open(output_file, "w") as output_stream:
            output_stream.write(output)
    else:
        print(output)


def validate_data(test_data):
    required_keys = {"nspkg", "ns"}
    missing_keys = required_keys - test_data.keys()
    if missing_keys:
        logger.error(f"Missing required keys in YAML: {', '.join(missing_keys)}")
        raise ValueError(f"Invalid test data structure: missing keys {missing_keys}")


def complete_data(test_data):
    nspkg_indexes = {nspkg["name"]: i + 1 for i, nspkg in enumerate(test_data["nspkg"])}
    ns_indexes = {ns["name"]: i + 1 for i, ns in enumerate(test_data["ns"])}
    for ns in test_data["ns"]:
        if ns["nspkg"] not in nspkg_indexes:
            logger.error(
                f"nspkg '{ns['nspkg']}' referenced in 'ns' not found in 'nspkg'"
            )
            raise ValueError(f"Invalid nspkg name: {ns['nspkg']}")
        ns["nspkg_index"] = nspkg_indexes[ns["nspkg"]]
    for vnf in test_data["vnf"]:
        if vnf["ns"] not in ns_indexes:
            logger.error(f"ns '{vnf['ns']}' referenced in 'vnf' not found in 'ns'")
            raise ValueError(f"Invalid ns name: {vnf['ns']}")
        vnf["ns_index"] = ns_indexes[vnf["ns"]]


####################################
# Main
####################################
if __name__ == "__main__":
    # Argument parse
    parser = argparse.ArgumentParser(
        description="Generate OSM tests from YAML configuration file."
    )
    parser.add_argument(
        "-v", "--verbose", action="count", default=0, help="increase output verbosity"
    )
    # parser.add_argument("--basedir", default=".", help="basedir for the test")
    parser.add_argument(
        "--config", required=True, help="yaml configuration file to create the test"
    )
    parser.add_argument(
        "--template",
        default="test_template.j2",
        help="template file for rendering the test (default: test_template.j2)",
    )
    parser.add_argument(
        "--output", default=None, help="output file (default: standard output)"
    )
    args = parser.parse_args()

    # Calculate paths relative to the script directory
    script_dir = os.path.dirname(os.path.abspath(__file__))
    config_path = os.path.abspath(args.config)
    template_path = os.path.join(script_dir, args.template)

    # Initialize logger
    set_logger(args.verbose)

    # Load test_data
    try:
        with open(config_path, "r") as yaml_stream:
            test_data = yaml.safe_load(yaml_stream)
    except FileNotFoundError:
        exit_with_error(f"Configuration file '{args.config}' not found.")
    except yaml.YAMLError as e:
        exit_with_error(f"Error parsing YAML file '{args.config}': {e}")
    validate_data(test_data)
    complete_data(test_data)
    logger.debug(f"Test data:\n{yaml.safe_dump(test_data)}")
    if not os.path.exists(template_path):
        exit_with_error(f"Template file '{template_path}' not found.")
    render_template(template_path, test_data, args.output)