infix/test/spec/generate_spec.py

#!/bin/env python3
import os
import ast
import argparse
import io
import sys
import re
import yaml

from pathlib import Path

import graphviz

def replace_image_tag(text, test_dir):
    """
    Convert images added in the description and replace with the required ifdefs to work
    generating the test specifcation as well.
    """

    pattern = r"image::(?P<image_name>\S+)\[(?P<label>[^\]]*)\]"

    def repl(match):
        image_name = match.group("image_name")
        label = match.group("label")

        return f"""ifdef::topdoc[]
image::../../{test_dir}/{image_name}[{label}]
endif::topdoc[]
ifndef::topdoc[]
ifdef::testgroup[]
image::{'/'.join(Path(test_dir).parts[3:])}/{image_name}[{label}]
endif::testgroup[]
ifndef::testgroup[]
image::{image_name}[{label}]
endif::testgroup[]
endif::topdoc[]"""

    return re.sub(pattern, repl, text)

class TestStepVisitor(ast.NodeVisitor):
    """
    A custom test step visitor to grab the test description (docstring)
    and the test steps test.step(.....) for the test case.

    """
    def __init__(self):
        self.test_steps=[]
        self.name = ""
        self.description = ""

    def visit_Module(self, node):
        """Extract docstring from a test."""
        docstring = ast.get_docstring(node)
        if docstring:
            lines = docstring.splitlines()

            newline_parsed=False

            for line in lines:
                if self.name == "":
                    self.name=line.strip()
                elif newline_parsed is False and line == "":
                    newline_parsed = True
                    continue # Skip mandatory newline
                else:
                    self.description=f"{self.description}\n{line}"
            self.description=self.description.strip()
        self.generic_visit(node)

    # Check for test.step() for the actual test steps
    def visit_Call(self, node):
        """Extract test.step from test"""
        if isinstance(node.func, ast.Attribute) and node.func.attr == 'step':
            if isinstance(node.func.value, ast.Name) and node.func.value.id == 'test':
                if node.args and isinstance(node.args[0], ast.Constant):
                    self.test_steps.append(node.args[0].value)
        self.generic_visit(node)  # Continue visiting other nodes

class TestCase:
    """All test specifcation resources for a test case"""
    def __init__(self, directory, rootdir=None):
        self.test_dir=Path(directory)
        if rootdir:
            rootdir=Path(f"{rootdir}")
            self.test_dir=self.test_dir.relative_to(rootdir)
        self.topology_dot=f"{directory}/topology.dot"
        self.topology_image=f"{directory}/topology"
        self.test_case=f"{directory}/test.py"
        self.specification=f"{directory}/Readme.adoc"

    def generate_topology(self):
        """Generate SVG file from the topology.dot file"""
        with open(self.topology_dot, 'r') as dot_file:
            dot_graph = dot_file.read()
            graph = graphviz.Source(dot_graph)
            graph.render(self.topology_image, format='svg', cleanup=True)
            pattern = r'<!--\s*Generated by graphviz.*?-->'
            content=""
            with open(f"{self.topology_image}.svg", "r") as f:
                content = f.read()
            mod_content = re.sub(pattern, '',content, flags=re.DOTALL)
            with open(f"{self.topology_image}.svg", 'w') as f:
                f.write(mod_content)


    def generate_specification(self, name, case_path, spec_path):
        """Generate a ASCIIDOC specification for the test case"""
        with open(case_path, 'r') as file:
            script_content = file.read()

        parsed_script = ast.parse(script_content)
        visitor = TestStepVisitor()
        visitor.visit(parsed_script)
        description=visitor.description if visitor.description != "" else "Undefined"
        test_steps=visitor.test_steps
        if name is None:
            name = visitor.name
        self.generate_topology()
        description = replace_image_tag(description, self.test_dir)
        with open(spec_path, "w") as spec:
            spec.write(f"=== {name}\n")
            spec.write("==== Description\n")
            spec.write(description + "\n\n")
            spec.write("==== Topology\n")
            spec.write("ifdef::topdoc[]\n")
            spec.write(f"image::../../{self.test_dir}/topology.svg[{name} topology]\n")
            spec.write("endif::topdoc[]\n")
            spec.write("ifndef::topdoc[]\n")
            spec.write("ifdef::testgroup[]\n")
            spec.write(f"image::{Path(*self.test_dir.parts[3:])}/topology.svg[{name} topology]\n")
            spec.write("endif::testgroup[]\n")
            spec.write("ifndef::testgroup[]\n")
            spec.write(f"image::topology.svg[{name} topology]\n")
            spec.write("endif::testgroup[]\n")
            spec.write("endif::topdoc[]\n")
            spec.write("==== Test sequence\n")
            spec.writelines([f". {step}\n" for step in test_steps])
            spec.write("\n\n<<<\n\n") # need empty lines to pagebreak

def parse_suite(directory, root, suitefile):
    test_spec = None
    readme = None
    num_tests = 0
    with open(suitefile, "r") as f:
        data = yaml.safe_load(f)

    for entry in data:
        if entry.get("suite"):
            path = os.path.dirname(entry["suite"])
            parse_suite(f"{directory}/{path}", root, f"{directory}/{entry['suite']}")
        elif entry.get("case"):
            if not entry.get("infamy"):
                test_title = None # Set from the docstring in the test case
            else:
                specification = entry["infamy"].get("specification", True);
                if not specification:
                    continue
                test_title = entry["infamy"]["title"]
            path = os.path.dirname(entry["case"])
            if readme is None:
                readme = open(f"{directory}/{path}/Readme.adoc.tmp", "w")
            test_case_id = entry["name"]
            test_file = f"{directory}/{entry['case']}"
            test_spec = f"{directory}/{path}/{test_case_id}.adoc"
            test_case = TestCase(f"{directory}/{path}", root)
            test_case.generate_specification(test_title, test_file, test_spec)
            readme.write(f"include::{path}{test_case_id}.adoc[]\n\n")
            if path != "":
                readme.close()
                readme = None
                os.unlink(f"{directory}/{path}/Readme.adoc")
                os.symlink(f"{test_case_id}.adoc", f"{directory}/{path}/Readme.adoc")
    if not readme is None:
        readme.close()
        os.rename(f"{directory}/Readme.adoc.tmp", f"{directory}/Readme.adoc")

parser = argparse.ArgumentParser(description="Generate a test specification for a subtree.")
parser.add_argument("-s", "--suite", required=True, help="The suite file to parse")
parser.add_argument("-r", "--root-dir", required=True, help="Path that all paths should be relative to")
args=parser.parse_args()

error_string = ""
output_capture = io.StringIO()
sys.stderr = output_capture

# This is hacky, graphviz output error only to stdout and return successful(always).
# If everything goes well, output shall be empty, fail on any output
output_capture.truncate(0)
output_capture.seek(0)
parse_suite(os.path.dirname(args.suite), args.root_dir, args.suite)

sys.stdout = sys.__stdout__

if len(error_string) > 0:
    print(error_string)
    exit(1)

exit(0)