xk6-frostfs/scenarios/preset/helpers/frostfs_cli.py

import re
from helpers.cmd import execute_cmd, log

def create_container(endpoint, policy, container_creation_retry, wallet_path, config, local=False, retry=0):
    if retry > int(container_creation_retry):
        raise ValueError(f"unable to create container: too many unsuccessful attempts")

    if wallet_path:
        wallet_file = f"--wallet {wallet_path}"
    if config:
        wallet_config = f"--config {config}"
    cmd_line = f"frostfs-cli --rpc-endpoint {endpoint} container create {wallet_file} {wallet_config} " \
               f" --policy '{policy}' --await"

    output, success = execute_cmd(cmd_line)

    if not success:
        log(f"{cmd_line}\n"
            f"Container has not been created\n"
            f"{output}", endpoint)
        return False
    
    try:
        fst_str = output.split('\n')[0]
    except Exception:
        log(f"{cmd_line}\n"
            f"Incorrect output\n"
            f"Output: {output or '<empty>'}", endpoint)
        return False
    splitted = fst_str.split(": ")
    if len(splitted) != 2:
        raise ValueError(f"no CID was parsed from command output:\t{fst_str}")
    cid = splitted[1]

    log(f"Created container: {cid} ({policy})", endpoint)

    if not local:
        return cid
    
    cmd_line = f"frostfs-cli netmap nodeinfo --rpc-endpoint {endpoint} {wallet_file} {wallet_config}"
    output, success = execute_cmd(cmd_line)

    if not success:
        log(f"{cmd_line}\n"
            f"Failed to get nodeinfo\n"
            f"{output}", endpoint)
        return False
    
    try:
        fst_str = output.split('\n')[0]
    except Exception:
        log(f"{cmd_line}\n"
            f"Incorrect output\n"
            f"Output: {output or '<empty>'}", endpoint)
        return False
    splitted = fst_str.split(": ")
    if len(splitted) != 2 or len(splitted[1]) == 0:
        raise ValueError(f"no node key was parsed from command output:\t{fst_str}")
    
    node_key = splitted[1]

    cmd_line = f"frostfs-cli container nodes --rpc-endpoint {endpoint} {wallet_file} {wallet_config} --cid {cid}"
    output, success = execute_cmd(cmd_line)

    if not success:
        log(f"{cmd_line}\n"
            f"Failed to get container nodes\n"
            f"{output}", endpoint)
        return False

    for output_str in output.split('\n'):
        output_str = output_str.lstrip().rstrip()
        if not output_str.startswith("Node "):
            continue
        splitted = output_str.split(": ")
        if len(splitted) != 2 or len(splitted[1]) == 0:
            continue
        try:
            k = splitted[1].split(" ")[0]
        except Exception:
            log(f"{cmd_line}\n"
                f"Incorrect output\n"
                f"Output: {output or '<empty>'}", endpoint)
            continue
        if k == node_key:
            return cid
    
    log(f"Created container {cid} is not stored on {endpoint}, creating another one...", endpoint)
    return create_container(endpoint, policy, container_creation_retry, wallet_path, config, local, retry + 1)


def upload_object(container, payload_filepath, endpoint, wallet_file, wallet_config):
    object_name = ""
    if wallet_file:
        wallet_file = "--wallet " + wallet_file
    if wallet_config:
        wallet_config = "--config " + wallet_config
    cmd_line = f"frostfs-cli --rpc-endpoint {endpoint} object put --file {payload_filepath} {wallet_file} {wallet_config} " \
               f"--cid {container} --no-progress"
    output, success = execute_cmd(cmd_line)

    if not success:
        log(f"{cmd_line}\n"
            f"Object {object_name} has not been uploaded\n"
            f"Error: {output}", endpoint)
        return False

    try:
        # taking second string from command output
        snd_str = output.split('\n')[1]
    except Exception:
        log(f"{cmd_line}\n"
            f"Incorrect output\n"
            f"Output: {output or '<empty>'}", endpoint)
        return False
    splitted = snd_str.split(": ")
    if len(splitted) != 2:
        raise Exception(f"no OID was parsed from command output: \t{snd_str}")
    return container, endpoint, splitted[1]


def get_object(cid, oid, endpoint, out_filepath, wallet_file, wallet_config):
    if wallet_file:
        wallet_file = "--wallet " + wallet_file
    if wallet_config:
        wallet_config = "--config " + wallet_config
    cmd_line = f"frostfs-cli object get -r {endpoint} --cid {cid} --oid {oid} {wallet_file} {wallet_config} " \
               f"--file {out_filepath}"

    output, success = execute_cmd(cmd_line)

    if not success:
        log(f"{cmd_line}\n"
            f"Failed to get object {oid} from container {cid}\n"
            f"Error: {output}", endpoint)
        return False

    return True


def search_object_by_id(cid, oid, endpoint, wallet_file, wallet_config, ttl=2):
    if wallet_file:
        wallet_file = "--wallet " + wallet_file
    if wallet_config:
        wallet_config = "--config " + wallet_config
    cmd_line = f"frostfs-cli object search --ttl {ttl} -r {endpoint} --cid {cid} --oid {oid} {wallet_file} {wallet_config} "

    output, success = execute_cmd(cmd_line)

    if not success:
        log(f"{cmd_line}\n"
            f"Failed to search object {oid} for container {cid}\n"
            f"Error: {output}", endpoint)
        return False

    re_rst = re.search(r'Found (\d+) objects', output)

    if not re_rst:
        raise Exception("Failed to parse search results")

    return re_rst.group(1)