utils.py

import os
import subprocess
import time
from pathlib import Path
from random import choice
from string import ascii_lowercase

from git import Repo
from run_iocsh import IOC


class Wrapper:
    def __init__(
        self, root_path: Path, *, name=None, include_dbd=True, url=None, **kwargs
    ):
        test_env = os.environ.copy()

        assert "EPICS_BASE" in test_env
        assert "E3_REQUIRE_VERSION" in test_env
        assert "EPICS_HOST_ARCH" in test_env

        self.epics_base = Path(os.getenv("EPICS_BASE"))
        self.base_version = self.epics_base.name.split("base-")[-1]
        self.host_arch = os.getenv("EPICS_HOST_ARCH")
        self.require_version = os.getenv("E3_REQUIRE_VERSION")

        e3_require_config = (
            self.epics_base / "require" / self.require_version / "configure"
        )

        assert e3_require_config.is_dir()

        if name is None:
            name = "test_mod_" + "".join(choice(ascii_lowercase) for _ in range(16))
        self.name = name
        self.version = "0.0.0+0"

        self.path = root_path / f"e3-{self.name}"
        self.config_dir = self.path / "configure"
        self.config_module = self.config_dir / "CONFIG_MODULE"
        module_path = (
            self.name
            if "E3_MODULE_SRC_PATH" not in kwargs
            else kwargs["E3_MODULE_SRC_PATH"]
        )
        self.module_dir = self.path / module_path
        self.build_dir = self.module_dir / f"O.{self.base_version}_{self.host_arch}"

        self.makefile = self.path / "Makefile"
        self.module_makefile = self.path / f"{self.name}.Makefile"

        self.url = url
        if self.url:
            Repo.clone_from(self.url, self.path)
        else:
            self.module_dir.mkdir(parents=True)
            self.config_dir.mkdir()
            self.config_module.touch()

            makefile_contents = f"""
TOP:=$(CURDIR)

E3_MODULE_NAME:={name}
E3_MODULE_VERSION:={self.version}
E3_MODULE_SRC_PATH:={module_path}
E3_MODULE_MAKEFILE:={name}.Makefile

include $(TOP)/configure/CONFIG_MODULE
-include $(TOP)/configure/CONFIG_MODULE.local

REQUIRE_CONFIG:={e3_require_config}

include $(REQUIRE_CONFIG)/CONFIG
include $(REQUIRE_CONFIG)/RULES_SITEMODS
"""
            self.makefile.write_text(makefile_contents)

            module_makefile_contents = """
where_am_I := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
include $(E3_REQUIRE_TOOLS)/driver.makefile
"""
            self.module_makefile.write_text(module_makefile_contents)

            if include_dbd:
                self.add_file("test.dbd")

            Repo.init(self.path)

    def add_file(self, name):
        (self.module_dir / name).touch()

    def write_dot_local_data(self, config_file: str, config_vars: dict):
        """Write config data to the specific .local file."""

        with open(self.config_dir / f"{config_file}.local", "w") as f:
            for var, value in config_vars.items():
                f.write(f"{var} = {value}\n")

    def add_var_to_config_module(self, makefile_var: str, value: str, modifier="+"):
        with open(self.config_module, "a") as f:
            f.write(f"export {makefile_var} {modifier}= {value}\n")

    def add_var_to_module_makefile(self, makefile_var: str, value: str, modifier="+"):
        with open(self.module_makefile, "a") as f:
            f.write(f"{makefile_var} {modifier}= {value}\n")

    def get_env_var(self, var: str):
        """Fetch an environment variable from the module build environment."""

        _, out, _ = self.run_make("cellvars")
        for line in out.split("\n"):
            if line.startswith(var):
                return line.split("=")[1].strip()
        return ""

    def run_make(
        self,
        target: str,
        *args,
        module_version: str = "",
        cell_path: str = "",
        **kwargs,
    ):
        """Attempt to run `make <target> <args>` for the current wrapper."""

        # We should not install in the global environment during tests
        if target == "install":
            target = "cellinstall"

        if target == "uninstall":
            target = "celluninstall"

        env = os.environ.copy()
        env.update(kwargs)

        args = list(args)
        if module_version:
            args.append(f"E3_MODULE_VERSION={module_version}")
        if not ("TEST_DEBUG_ARCH" in env and not env["TEST_DEBUG_ARCH"]):
            args.append("EXCLUDE_ARCHS=debug")

        if cell_path:
            self.write_dot_local_data("CONFIG_CELL", {"E3_CELL_PATH": cell_path})
        make_cmd = ["make", "-C", self.path, target] + args
        p = subprocess.Popen(
            make_cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            encoding="utf-8",
            env=env,
        )
        outs, errs = p.communicate()
        return p.returncode, outs, errs


def run_ioc_get_output(*args, **kwargs):
    """Run an IOC and try to load the test module."""
    ioc_args = []
    module = kwargs.get("module", None)
    version = kwargs.get("version", None)
    if module:
        ioc_args.append("-r")
        ioc_args.append(f"{module},{version}" if version else module)
    cell_path = kwargs.get("cell_path", None)
    if cell_path:
        ioc_args.extend(["-l", cell_path])
    ioc_args.extend(args)
    with IOC(*ioc_args, ioc_executable="iocsh") as ioc:
        time.sleep(1)
    return ioc.proc.returncode, ioc.outs, ioc.errs