frostfs-testcases/pytest_tests/helpers/container_creation.py

import json
import time

from frostfs_testlib import reporter
from frostfs_testlib.cli.frostfs_cli.cli import FrostfsCli
from frostfs_testlib.resources.common import MORPH_BLOCK_TIME
from frostfs_testlib.shell.interfaces import Shell
from frostfs_testlib.steps.cli.container import create_container, search_nodes_with_container
from frostfs_testlib.storage.cluster import Cluster
from frostfs_testlib.storage.dataclasses import ape
from frostfs_testlib.storage.dataclasses.wallet import WalletInfo
from frostfs_testlib.testing.parallel import parallel
from frostfs_testlib.utils import datetime_utils

from .container_request import ContainerRequest, MultipleContainersRequest


def create_container_with_ape(
    container_request: ContainerRequest,
    frostfs_cli: FrostfsCli,
    wallet: WalletInfo,
    shell: Shell,
    cluster: Cluster,
    endpoint: str,
) -> str:
    with reporter.step("Create container"):
        cid = _create_container_by_spec(container_request, wallet, shell, cluster, endpoint)

    if container_request.ape_rules:
        with reporter.step("Apply APE rules for container"):
            _apply_ape_rules(cid, frostfs_cli, endpoint, container_request.ape_rules)

        with reporter.step("Wait for one block"):
            time.sleep(datetime_utils.parse_time(MORPH_BLOCK_TIME))

    with reporter.step("Search nodes holding the container"):
        container_holder_nodes = search_nodes_with_container(wallet, cid, shell, cluster.default_rpc_endpoint, cluster)
        report_data = {node.id: node.host_ip for node in container_holder_nodes}

    reporter.attach(json.dumps(report_data, indent=2), "container_nodes.json")

    return cid


@reporter.step("Create multiple containers with APE")
def create_containers_with_ape(
    frostfs_cli: FrostfsCli,
    wallet: WalletInfo,
    shell: Shell,
    cluster: Cluster,
    endpoint: str,
    multiple_containers_request: MultipleContainersRequest,
) -> list[str]:
    cids_futures = parallel(create_container_with_ape, multiple_containers_request, frostfs_cli, wallet, shell, cluster, endpoint)
    return [future.result() for future in cids_futures]


@reporter.step("Create container by spec {container_request}")
def _create_container_by_spec(
    container_request: ContainerRequest,
    wallet: WalletInfo,
    shell: Shell,
    cluster: Cluster,
    endpoint: str,
) -> str:
    return create_container(
        wallet,
        shell,
        endpoint,
        container_request.parsed_rule(cluster),
        wait_for_creation=False,
        nns_zone=container_request.ns_zone,
        nns_name=container_request.ns_name,
    )


def _apply_ape_rules(cid: str, frostfs_cli: FrostfsCli, endpoint: str, ape_rules: list[ape.Rule]):
    for ape_rule in ape_rules:
        rule_str = ape_rule.as_string()
        with reporter.step(f"Apply APE rule '{rule_str}' for container {cid}"):
            frostfs_cli.ape_manager.add(
                endpoint,
                ape_rule.chain_id,
                target_name=cid,
                target_type="container",
                rule=rule_str,
            )