test.py

"""The test module helps writing tests protogen protoc plugins."""

from importlib import import_module
from typing import Dict, List, Tuple
import io
import os.path
import subprocess
import sys
import tempfile

import protogen

import google.protobuf.compiler.plugin_pb2
import google.protobuf.descriptor_pb2


def run_plugin(
    proto_paths: List[str],
    files_to_generate: List[str],
    plugin: str,
    additional_protoc_args: List[str] = ["--experimental_allow_proto3_optional"],
    parameter: Dict[str, str] = {},
) -> protogen.CodeGeneratorResponse:
    """Run a protoc plugin python module.

    Runs the protoc plugin module `plugin` within the current process (as
    opposed to protoc which calls the plugin in a subprocess). This makes it
    easy to attach debugger to the current process to debug the plugin.

    Arguments
    ---------
    proto_paths : List[str]
        List of directories that act as proto paths (will be passed via
        --proto_path/-I flag to protoc).
    files_to_generate : List[str]
        List of proto files to generate output for (will be passed as positional
        arguments to protoc).
    plugin : str
        Python module name of the plugin to run. The plugin will be called and
        must be loadable and executable via `importlib.import_module(plugin)`.
    additional_protc_args : List[str]
        Additional arguments that will be passed to protoc.
    parameter : Dict[str, str]
        Parameter that will be passed to the plugin.

    Returns
    -------
    response: protogen.CodeGeneratorResponse
        The response generated by the plugin.
    """

    req = _prepare_code_generator_request(
        proto_paths, files_to_generate, additional_protoc_args, parameter
    )

    # Open stdin and stdout for the plugin.
    # We need SpoledTemporaryFile(mode="w+t")._file because its a TextIOWrapper
    # which sys.stdin and sys.stdout also are.
    with tempfile.SpooledTemporaryFile(mode="w+t") as tmp_stdin:
        with tempfile.SpooledTemporaryFile(mode="w+t") as tmp_stdout:
            fake_stdin = tmp_stdin._file
            fake_stdin.buffer.write(req.SerializeToString())
            fake_stdin.flush()
            fake_stdin.seek(0)

            fake_stdout = tmp_stdout._file

            _stdin, sys.stdin = sys.stdin, fake_stdin
            _stdout, sys.stdout = sys.stdout, fake_stdout

            # Call the plugin under test.
            import_module(plugin)

            fake_stdout.seek(0)
            resp = google.protobuf.compiler.plugin_pb2.CodeGeneratorResponse.FromString(
                fake_stdout.buffer.read()
            )

    # Reset stdin and stdout.
    sys.stdin = _stdin
    sys.stdout = _stdout

    return protogen.CodeGeneratorResponse(resp)


def _prepare_code_generator_request(
    proto_paths: List[str],
    files_to_generate: List[str],
    additional_protoc_args: List[str],
    parameter: Dict[str, str],
):
    req = google.protobuf.compiler.plugin_pb2.CodeGeneratorRequest(
        file_to_generate=files_to_generate,
        parameter=",".join([f"{k}={v}" for (k, v) in parameter.items()]),
        proto_file=[],
        compiler_version=google.protobuf.compiler.plugin_pb2.Version(
            major=1, minor=2, patch=3, suffix=""
        ),
    )

    with tempfile.TemporaryDirectory() as tmpdirname:
        f = os.path.join(tmpdirname, "descriptor_set.pb")

        cmd = ["protoc"]
        for proto_path in proto_paths:
            cmd.extend(["-I", proto_path])
        cmd.append(f"--descriptor_set_out={f}")
        cmd.append("--include_imports")
        cmd.append("--include_source_info")
        cmd.extend(additional_protoc_args)
        cmd.extend(files_to_generate)

        code, output = _run_protoc(cmd)
        if code != 0:
            raise Exception(output)

        ff = io.open(f, "rb")
        desc_set = google.protobuf.descriptor_pb2.FileDescriptorSet.FromString(
            ff.read()
        )
        req.proto_file.extend(desc_set.file)

    return req


def _run_protoc(args: List[str]) -> Tuple[int, str]:
    proc = subprocess.Popen(
        args, text=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE
    )
    # Executed from current directory (repo root)
    code = proc.wait()
    if code == 0:
        output = proc.stdout.read()
    else:
        output = proc.stderr.read()
    return code, output