pluginbase.py

import logging
import threading
import typing

from abc import ABC
from abc import abstractmethod
from typing import Optional

from ktoolbox import common

from tftbase import PluginOutput
from tftbase import PluginResult
from tftbase import TestMetadata


logger = logging.getLogger("tft." + __name__)


class Plugin(ABC):
    PLUGIN_NAME = ""

    @property
    def log_name(self) -> str:
        return f"plugin[{self.PLUGIN_NAME}"

    def enable(
        self,
        *,
        ts: "TestSettings",
        node_server_name: str,
        node_client_name: str,
        perf_server: "ServerTask",
        perf_client: "ClientTask",
        tenant: bool,
    ) -> list["PluginTask"]:
        tasks = self._enable(
            ts=ts,
            node_server_name=node_server_name,
            node_client_name=node_client_name,
            perf_server=perf_server,
            perf_client=perf_client,
            tenant=tenant,
        )
        logger.debug(f"{self.log_name}: enable ({len(tasks)} tasks, {tasks})")
        return tasks

    @abstractmethod
    def _enable(
        self,
        *,
        ts: "TestSettings",
        node_server_name: str,
        node_client_name: str,
        perf_server: "ServerTask",
        perf_client: "ClientTask",
        tenant: bool,
    ) -> list["PluginTask"]:
        pass

    def eval_plugin_output(
        self,
        md: TestMetadata,
        plugin_output: PluginOutput,
    ) -> Optional[PluginResult]:
        if not plugin_output.success:
            logger.error(
                f"{self.PLUGIN_NAME} plugin failed for {common.dataclass_to_json(md)}: {plugin_output.err_msg}"
            )
        else:
            logger.debug(
                f"{self.PLUGIN_NAME} plugin succeded for {common.dataclass_to_json(md)}"
            )
        return PluginResult(
            tft_metadata=md,
            plugin_name=self.PLUGIN_NAME,
            success=plugin_output.success,
            msg=plugin_output.err_msg,
            plugin_output=plugin_output,
        )


_plugin_registry_lock = threading.Lock()

_plugin_registry: dict[str, Plugin] = {}


def _get_plugin_registry() -> dict[str, Plugin]:
    # Plugins self-register via pluginbase.register_plugin()
    # when being imported. Ensure they are imported.
    import pluginMeasureCpu  # noqa: F401
    import pluginMeasurePower  # noqa: F401
    import pluginValidateOffload  # noqa: F401

    return _plugin_registry


def register_plugin(plugin: Plugin) -> Plugin:
    name = plugin.PLUGIN_NAME
    with _plugin_registry_lock:
        p2 = _plugin_registry.setdefault(name, plugin)
    if p2 is not plugin:
        raise ValueError(f"plugin {repr(name)} is already registered")
    return plugin


def get_all() -> list[Plugin]:
    registry = _get_plugin_registry()
    with _plugin_registry_lock:
        return [registry[k] for k in sorted(registry)]


def get_by_name(plugin_name: str) -> Plugin:
    registry = _get_plugin_registry()
    with _plugin_registry_lock:
        plugin = registry.get(plugin_name)
    if plugin is None:
        raise ValueError(f'Plugin "{plugin_name}" does not exist')
    return plugin


if typing.TYPE_CHECKING:
    # "pluginbase" cannot import modules like perf, task or testSettings, because
    # those modules import "pluginbase" in turn. However, to forward declare
    # type annotations, we do need those module here. Import them with
    # TYPE_CHECKING, but otherwise avoid the cyclic dependency between
    # modules.
    from task import ClientTask
    from task import ServerTask
    from task import PluginTask
    from testSettings import TestSettings