frostfs-testcases/pytest_tests/helpers/cluster.py

import pathlib
import random
import re
from dataclasses import dataclass
from typing import Any, Optional

import pytest
import yaml
from frostfs_testlib.blockchain import RPCClient
from frostfs_testlib.hosting import Host, Hosting
from frostfs_testlib.hosting.config import ServiceConfig
from frostfs_testlib.utils import wallet_utils


@dataclass
class NodeBase:
    """
    Represents a node of some underlying service
    """

    id: str
    name: str
    host: Host

    def __init__(self, id, name, host) -> None:
        self.id = id
        self.name = name
        self.host = host
        self.construct()

    def construct(self):
        pass

    def __eq__(self, other):
        return self.name == other.name

    def __hash__(self):
        return id(self.name)

    def __str__(self):
        return self.label

    def __repr__(self) -> str:
        return self.label

    @property
    def label(self) -> str:
        return self.name

    def start_service(self):
        self.host.start_service(self.name)

    def stop_service(self):
        self.host.stop_service(self.name)

    def get_wallet_password(self) -> str:
        return self._get_attribute(_ConfigAttributes.WALLET_PASSWORD)

    def get_wallet_path(self) -> str:
        return self._get_attribute(
            _ConfigAttributes.LOCAL_WALLET_PATH,
            _ConfigAttributes.WALLET_PATH,
        )

    def get_remote_wallet_path(self) -> str:
        """
        Returns node wallet file path located on remote host
        """
        return self._get_attribute(
            _ConfigAttributes.WALLET_PATH,
        )

    def get_remote_config_path(self) -> str:
        """
        Returns node config file path located on remote host
        """
        return self._get_attribute(
            _ConfigAttributes.CONFIG_PATH,
        )

    def get_wallet_config_path(self):
        return self._get_attribute(
            _ConfigAttributes.LOCAL_WALLET_CONFIG,
            _ConfigAttributes.WALLET_CONFIG,
        )

    def get_wallet_public_key(self):
        storage_wallet_path = self.get_wallet_path()
        storage_wallet_pass = self.get_wallet_password()
        return wallet_utils.get_wallet_public_key(storage_wallet_path, storage_wallet_pass)

    def _get_attribute(self, attribute_name: str, default_attribute_name: str = None) -> list[str]:
        config = self.host.get_service_config(self.name)
        if default_attribute_name:
            return config.attributes.get(
                attribute_name, config.attributes.get(default_attribute_name)
            )
        else:
            return config.attributes.get(attribute_name)

    def _get_service_config(self) -> ServiceConfig:
        return self.host.get_service_config(self.name)


class InnerRingNode(NodeBase):
    """
    Class represents inner ring node in a cluster

    Inner ring node is not always the same as physical host (or physical node, if you will):
        It can be service running in a container or on physical host
    For testing perspective, it's not relevant how it is actually running,
        since frostfs network will still treat it as "node"
    """

    def get_netmap_cleaner_threshold(self) -> int:
        config_file = self.get_remote_config_path()
        contents = self.host.get_shell().exec(f"cat {config_file}").stdout

        config = yaml.safe_load(contents)
        value = config["netmap_cleaner"]["threshold"]

        return int(value)


class S3Gate(NodeBase):
    """
    Class represents S3 gateway in a cluster
    """

    def get_endpoint(self) -> str:
        return self._get_attribute(_ConfigAttributes.ENDPOINT_DATA)

    @property
    def label(self) -> str:
        return f"{self.name}: {self.get_endpoint()}"


class HTTPGate(NodeBase):
    """
    Class represents HTTP gateway in a cluster
    """

    def get_endpoint(self) -> str:
        return self._get_attribute(_ConfigAttributes.ENDPOINT_DATA)

    @property
    def label(self) -> str:
        return f"{self.name}: {self.get_endpoint()}"


class MorphChain(NodeBase):
    """
    Class represents side-chain aka morph-chain consensus node in a cluster

    Consensus node is not always the same as physical host (or physical node, if you will):
        It can be service running in a container or on physical host
    For testing perspective, it's not relevant how it is actually running,
        since frostfs network will still treat it as "node"
    """

    rpc_client: RPCClient = None

    def construct(self):
        self.rpc_client = RPCClient(self.get_endpoint())

    def get_endpoint(self) -> str:
        return self._get_attribute(_ConfigAttributes.ENDPOINT_INTERNAL)

    @property
    def label(self) -> str:
        return f"{self.name}: {self.get_endpoint()}"


class MainChain(NodeBase):
    """
    Class represents main-chain consensus node in a cluster

    Consensus node is not always the same as physical host:
        It can be service running in a container or on physical host (or physical node, if you will):
    For testing perspective, it's not relevant how it is actually running,
        since frostfs network will still treat it as "node"
    """

    rpc_client: RPCClient = None

    def construct(self):
        self.rpc_client = RPCClient(self.get_endpoint())

    def get_endpoint(self) -> str:
        return self._get_attribute(_ConfigAttributes.ENDPOINT_INTERNAL)

    @property
    def label(self) -> str:
        return f"{self.name}: {self.get_endpoint()}"


class StorageNode(NodeBase):
    """
    Class represents storage node in a storage cluster

    Storage node is not always the same as physical host:
        It can be service running in a container or on physical host (or physical node, if you will):
    For testing perspective, it's not relevant how it is actually running,
        since frostfs network will still treat it as "node"
    """

    def get_rpc_endpoint(self) -> str:
        return self._get_attribute(_ConfigAttributes.ENDPOINT_DATA)

    def get_control_endpoint(self) -> str:
        return self._get_attribute(_ConfigAttributes.CONTROL_ENDPOINT)

    def get_un_locode(self):
        return self._get_attribute(_ConfigAttributes.UN_LOCODE)

    @property
    def label(self) -> str:
        return f"{self.name}: {self.get_rpc_endpoint()}"


class ClusterNode:
    """
    Represents physical node where multiple different services may be located
    """

    host: Host
    storage_node: Optional[StorageNode] = None
    ir_node: Optional[InnerRingNode] = None
    s3_gate: Optional[S3Gate] = None
    http_gate: Optional[HTTPGate] = None
    morph_chain: Optional[MorphChain] = None
    main_chain: Optional[MainChain] = None

    def __init__(self, host: Host, nodes: list[NodeBase]) -> None:
        self.host = host
        attributes_map = {
            StorageNode: "storage_node",
            InnerRingNode: "ir_node",
            S3Gate: "s3_gate",
            HTTPGate: "http_gate",
            MorphChain: "morph_chain",
            MainChain: "main_chain",
        }

        for node in nodes:
            if node.host.config.address == host.config.address:
                self.__setattr__(attributes_map[node.__class__], node)

    @property
    def host_ip(self):
        return self.host.config.address

    def __eq__(self, other):
        return self.host.config.address == other.host.config.address

    def __hash__(self):
        return id(self.host.config.address)

    def __str__(self):
        return self.host.config.address

    def __repr__(self) -> str:
        return self.host.config.address

    def get_service_by_type(self, service_type: type[NodeBase]) -> type[NodeBase]:
        class_name = service_type.__name__
        class_field_map = {
            StorageNode.__name__: self.storage_node,
            InnerRingNode.__name__: self.ir_node,
            S3Gate.__name__: self.s3_gate,
            HTTPGate.__name__: self.http_gate,
            MorphChain.__name__: self.morph_chain,
        }
        if class_name not in class_field_map:
            raise pytest.fail(f"Invalid type passed {class_name}")
        return class_field_map[class_name]

    def get_list_of_services(self) -> list[str]:
        return [
            self.storage_node.get_service_systemctl_name(),
            self.ir_node.get_service_systemctl_name(),
            self.s3_gate.get_service_systemctl_name(),
            self.http_gate.get_service_systemctl_name(),
            self.morph_chain.get_service_systemctl_name(),
        ]


class Cluster:
    """
    This class represents a Cluster object for the whole storage based on provided hosting
    """

    default_rpc_endpoint: str
    default_s3_gate_endpoint: str
    default_http_gate_endpoint: str
    cluster_nodes: list[ClusterNode]

    def __init__(self, hosting: Hosting) -> None:
        self._hosting = hosting
        self.default_rpc_endpoint = self.storage_nodes[0].get_rpc_endpoint()
        self.default_s3_gate_endpoint = self.s3gates[0].get_endpoint()
        self.default_http_gate_endpoint = self.http_gates[0].get_endpoint()

    @property
    def hosts(self) -> list[Host]:
        """
        Returns list of Hosts
        """
        return self._hosting.hosts

    @property
    def cluster_nodes(self) -> list[ClusterNode]:
        """
        Returns list of Cluster Nodes
        """

        return [
            ClusterNode(host, self._find_nodes_by_pattern(f".*{id:02d}$"))
            for id, host in enumerate(self.hosts, start=1)
        ]

    @property
    def hosting(self) -> Hosting:
        return self._hosting

    def _create_wallet_config(self, service: ServiceConfig) -> None:
        wallet_path = service.attributes[_ConfigAttributes.LOCAL_WALLET_CONFIG]
        wallet_password = service.attributes[_ConfigAttributes.WALLET_PASSWORD]
        with open(wallet_path, "w") as file:
            yaml.dump({"password": wallet_password}, file)

    def create_wallet_configs(self, hosting: Hosting) -> None:
        configs = hosting.find_service_configs(".*")
        for config in configs:
            if _ConfigAttributes.LOCAL_WALLET_CONFIG in config.attributes:
                self._create_wallet_config(config)

    def is_local_devevn(self) -> bool:
        if len(self.hosting.hosts) == 1:
            host = self.hosting.hosts[0]
            if host.config.address == "localhost" and host.config.plugin_name == "docker":
                return True
        return False

    @property
    def storage_nodes(self) -> list[StorageNode]:
        """
        Returns list of Storage Nodes (not physical nodes)
        """
        return self._get_nodes(_ServicesNames.STORAGE)

    @property
    def s3gates(self) -> list[S3Gate]:
        """
        Returns list of S3 gates
        """
        return self._get_nodes(_ServicesNames.S3_GATE)

    @property
    def http_gates(self) -> list[S3Gate]:
        """
        Returns list of HTTP gates
        """
        return self._get_nodes(_ServicesNames.HTTP_GATE)

    @property
    def morph_chain_nodes(self) -> list[MorphChain]:
        """
        Returns list of morph-chain consensus nodes (not physical nodes)
        """
        return self._get_nodes(_ServicesNames.MORPH_CHAIN)

    @property
    def main_chain_nodes(self) -> list[MainChain]:
        """
        Returns list of main-chain consensus nodes (not physical nodes)
        """
        return self._get_nodes(_ServicesNames.MAIN_CHAIN)

    @property
    def ir_nodes(self) -> list[InnerRingNode]:
        """
        Returns list of inner-ring nodes (not physical nodes)
        """
        return self._get_nodes(_ServicesNames.INNER_RING)

    def _get_nodes(self, service_name) -> list[NodeBase]:
        return self._find_nodes_by_pattern(f"{service_name}\d*$")

    def _find_nodes_by_pattern(self, pattern) -> list[NodeBase]:
        configs = self.hosting.find_service_configs(pattern)

        class_mapping: dict[str, Any] = {
            _ServicesNames.STORAGE: StorageNode,
            _ServicesNames.INNER_RING: InnerRingNode,
            _ServicesNames.MORPH_CHAIN: MorphChain,
            _ServicesNames.S3_GATE: S3Gate,
            _ServicesNames.HTTP_GATE: HTTPGate,
            _ServicesNames.MAIN_CHAIN: MainChain,
        }

        found_nodes = []
        for config in configs:
            # config.name is something like s3-gate01. Cut last digits to know service type
            service_type = re.findall(".*\D", config.name)[0]
            # exclude unsupported services
            if service_type not in class_mapping.keys():
                continue

            cls = class_mapping.get(service_type)
            found_nodes.append(
                cls(
                    self._get_id(config.name),
                    config.name,
                    self.hosting.get_host_by_service(config.name),
                )
            )
        return found_nodes

    def _get_id(self, node_name) -> str:
        pattern = "\d*$"

        matches = re.search(pattern, node_name)
        if matches:
            return int(matches.group())

    def get_random_storage_rpc_endpoint(self) -> str:
        return random.choice(self.get_storage_rpc_endpoints())

    def get_random_storage_rpc_endpoint_mgmt(self) -> str:
        return random.choice(self.get_storage_rpc_endpoints_mgmt())

    def get_storage_rpc_endpoints(self) -> list[str]:
        nodes = self.storage_nodes
        return [node.get_rpc_endpoint() for node in nodes]

    def get_storage_rpc_endpoints_mgmt(self) -> list[str]:
        nodes = self.storage_nodes
        return [node.get_rpc_endpoint_mgmt() for node in nodes]

    def get_morph_endpoints(self) -> list[str]:
        nodes = self.morph_chain_nodes
        return [node.get_endpoint() for node in nodes]


class _ServicesNames:
    STORAGE = "s"
    S3_GATE = "s3-gate"
    HTTP_GATE = "http-gate"
    MORPH_CHAIN = "morph-chain"
    INNER_RING = "ir"
    MAIN_CHAIN = "main-chain"


class _ConfigAttributes:
    WALLET_PASSWORD = "wallet_password"
    WALLET_PATH = "wallet_path"
    WALLET_CONFIG = "wallet_config"
    CONFIG_PATH = "config_path"
    LOCAL_WALLET_PATH = "local_wallet_path"
    LOCAL_WALLET_CONFIG = "local_config_path"
    ENDPOINT_DATA = "endpoint_data0"
    ENDPOINT_INTERNAL = "endpoint_internal0"
    CONTROL_ENDPOINT = "control_endpoint"
    UN_LOCODE = "un_locode"