ClickHouse/tests/integration/test_disks_app_interactive/test.py

import io
import os
import pathlib
import select
import subprocess
from typing import Dict, List, Optional, Tuple, Union

import pytest

from helpers.cluster import ClickHouseCluster


class ClickHouseDisksException(Exception):
    pass


@pytest.fixture(scope="module")
def started_cluster():
    global cluster
    try:
        cluster = ClickHouseCluster(__file__)
        cluster.add_instance(
            "disks_app_test",
            main_configs=["server_configs/config.xml"],
            with_minio=True,
        )

        cluster.start()
        yield cluster

    finally:
        cluster.shutdown()


class DisksClient(object):
    SEPARATOR = b"\a\a\a\a\n"
    local_client: Optional["DisksClient"] = None  # static variable
    default_disk_root_directory: str = "/var/lib/clickhouse"

    def __init__(self, bin_path: str, config_path: str, working_path: str):
        self.bin_path = bin_path
        self.working_path = working_path

        self.proc = subprocess.Popen(
            [bin_path, "disks", "--test-mode", "--config", config_path],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        self.poller = select.epoll()
        self.poller.register(self.proc.stdout)
        self.poller.register(self.proc.stderr)

        self.stopped = False

        self._fd_nums = {
            self.proc.stdout.fileno(): self.proc.stdout,
            self.proc.stderr.fileno(): self.proc.stderr,
        }

    def execute_query(self, query: str, timeout: float = 5.0) -> str:
        output = io.BytesIO()

        self.proc.stdin.write(query.encode() + b"\n")
        self.proc.stdin.flush()

        events = self.poller.poll(timeout)
        if not events:
            raise TimeoutError(f"Disks client returned no output")

        for fd_num, event in events:
            if event & (select.EPOLLIN | select.EPOLLPRI):
                file = self._fd_nums[fd_num]

                if file == self.proc.stdout:
                    while True:
                        chunk = file.readline()
                        if chunk.endswith(self.SEPARATOR):
                            break

                        output.write(chunk)

                elif file == self.proc.stderr:
                    error_line = self.proc.stderr.readline()
                    print(error_line)
                    raise ClickHouseDisksException(error_line.strip().decode())

            else:
                raise ValueError(f"Failed to read from pipe. Flag {event}")

        data = output.getvalue().strip().decode()
        return data

    def list_disks(self) -> List[Tuple[str, str]]:
        output = self.execute_query("list-disks")
        return list(
            sorted(
                map(
                    lambda x: (x.split(":")[0], ":".join(x.split(":")[1:])),
                    output.split("\n"),
                )
            )
        )

    def current_disk_with_path(self) -> Tuple[str, str]:
        output = self.execute_query("current_disk_with_path")
        disk_line = output.split("\n")[0]
        path_line = output.split("\n")[1]
        assert disk_line.startswith("Disk: ")
        assert path_line.startswith("Path: ")
        return disk_line[6:], path_line[6:]

    def ls(
        self, path: str, recursive: bool = False, show_hidden: bool = False
    ) -> Union[List[str], Dict[str, List[str]]]:
        recursive_adding = "--recursive " if recursive else ""
        show_hidden_adding = "--all " if show_hidden else ""
        output = self.execute_query(
            f"list {path} {recursive_adding} {show_hidden_adding}"
        )
        if recursive:
            answer: Dict[str, List[str]] = dict()
            blocks = output.split("\n\n")
            for block in blocks:
                directory = block.split("\n")[0][:-1]
                files = block.split("\n")[1:]
                answer[directory] = files
            return answer
        else:
            return output.split("\n")

    def switch_disk(self, disk: str, directory: Optional[str] = None):
        directory_addition = f"--path {directory} " if directory is not None else ""
        self.execute_query(f"switch-disk {disk} {directory_addition}")

    def cd(self, directory: str, disk: Optional[str] = None):
        disk_addition = f"--disk {disk} " if disk is not None else ""
        self.execute_query(f"cd {directory} {disk_addition}")

    def copy(
        self,
        path_from,
        path_to,
        disk_from: Optional[str] = None,
        disk_to: Optional[str] = None,
        recursive: bool = False,
    ):
        disk_from_option = f"--disk-from {disk_from} " if disk_from is not None else ""
        disk_to_option = f"--disk-to {disk_to} " if disk_to is not None else ""
        recursive_tag = "--recursive" if recursive else ""

        self.execute_query(
            f"copy {recursive_tag} {path_from} {path_to} {disk_from_option} {disk_to_option}"
        )

    def move(self, path_from: str, path_to: str):
        self.execute_query(f"move {path_from} {path_to}")

    def rm(self, path: str, recursive: bool = False):
        recursive_tag = "--recursive" if recursive else ""
        self.execute_query(f"rm {recursive_tag} {path}")

    def mkdir(self, path: str, recursive: bool = False):
        recursive_adding = "--recursive " if recursive else ""
        self.execute_query(f"mkdir {path} {recursive_adding}")

    def ln(self, path_from: str, path_to: str):
        self.execute_query(f"link {path_from} {path_to}")

    def read(self, path_from: str, path_to: Optional[str] = None):
        path_to_adding = f"--path-to {path_to} " if path_to is not None else ""
        output = self.execute_query(f"read {path_from} {path_to_adding}")
        return output

    def write(
        self, path_from: str, path_to: str
    ):  # Writing from stdin is difficult to test (do not know how to do this in python)
        path_from_adding = f"--path-from {path_from}"
        self.execute_query(f"write {path_from_adding} {path_to}")

    @staticmethod
    def getLocalDisksClient(refresh: bool):
        if (DisksClient.local_client is None) or refresh:
            binary_file = os.environ.get("CLICKHOUSE_TESTS_SERVER_BIN_PATH")
            current_working_directory = str(pathlib.Path().resolve())
            config_file = f"{current_working_directory}/test_disks_app_interactive/configs/config.xml"
            if not os.path.exists(DisksClient.default_disk_root_directory):
                os.mkdir(DisksClient.default_disk_root_directory)

            DisksClient.local_client = DisksClient(
                binary_file, config_file, current_working_directory
            )
            return DisksClient.local_client
        else:
            return DisksClient.local_client


def test_disks_app_interactive_list_disks():
    client = DisksClient.getLocalDisksClient(True)
    expected_disks_with_path = [
        ("default", "/"),
        ("local", client.working_path),
    ]
    assert expected_disks_with_path == client.list_disks()
    assert client.current_disk_with_path() == ("default", "/")
    client.switch_disk("local")
    assert client.current_disk_with_path() == (
        "local",
        client.working_path,
    )


def test_disks_app_interactive_list_files_local():
    client = DisksClient.getLocalDisksClient(True)
    client.switch_disk("local")
    excepted_listed_files = sorted(os.listdir("test_disks_app_interactive/"))
    listed_files = sorted(client.ls("test_disks_app_interactive/"))
    assert excepted_listed_files == listed_files


def test_disks_app_interactive_list_directories_default():
    client = DisksClient.getLocalDisksClient(True)
    traversed_dir = client.ls(".", recursive=True)
    client.mkdir("dir1")
    client.mkdir("dir2")
    client.mkdir(".dir3")
    client.cd("dir1")
    client.mkdir("dir11")
    client.mkdir(".dir12")
    client.mkdir("dir13")
    client.cd("../dir2")
    client.mkdir("dir21")
    client.mkdir("dir22")
    client.mkdir(".dir23")
    client.cd("../.dir3")
    client.mkdir("dir31")
    client.mkdir(".dir32")
    client.cd("..")
    traversed_dir = client.ls(".", recursive=True)
    assert traversed_dir == {
        ".": ["dir1", "dir2"],
        "./dir1": ["dir11", "dir13"],
        "./dir2": ["dir21", "dir22"],
        "./dir1/dir11": [],
        "./dir1/dir13": [],
        "./dir2/dir21": [],
        "./dir2/dir22": [],
    }
    traversed_dir = client.ls(".", recursive=True, show_hidden=True)
    assert traversed_dir == {
        ".": [".dir3", "dir1", "dir2"],
        "./dir1": [".dir12", "dir11", "dir13"],
        "./dir2": [".dir23", "dir21", "dir22"],
        "./.dir3": [".dir32", "dir31"],
        "./dir1/dir11": [],
        "./dir1/.dir12": [],
        "./dir1/dir13": [],
        "./dir2/dir21": [],
        "./dir2/dir22": [],
        "./dir2/.dir23": [],
        "./.dir3/dir31": [],
        "./.dir3/.dir32": [],
    }
    client.rm("dir2", recursive=True)
    traversed_dir = client.ls(".", recursive=True, show_hidden=True)
    assert traversed_dir == {
        ".": [".dir3", "dir1"],
        "./dir1": [".dir12", "dir11", "dir13"],
        "./.dir3": [".dir32", "dir31"],
        "./dir1/dir11": [],
        "./dir1/.dir12": [],
        "./dir1/dir13": [],
        "./.dir3/dir31": [],
        "./.dir3/.dir32": [],
    }
    traversed_dir = client.ls(".", recursive=True, show_hidden=False)
    assert traversed_dir == {
        ".": ["dir1"],
        "./dir1": ["dir11", "dir13"],
        "./dir1/dir11": [],
        "./dir1/dir13": [],
    }
    client.rm("dir1", recursive=True)
    client.rm(".dir3", recursive=True)
    assert client.ls(".", recursive=True, show_hidden=False) == {".": []}


def test_disks_app_interactive_cp_and_read():
    initial_text = "File content"
    with open("a.txt", "w") as file:
        file.write(initial_text)
    client = DisksClient.getLocalDisksClient(True)
    client.switch_disk("default")
    client.copy("a.txt", "/a.txt", disk_from="local", disk_to="default")
    read_text = client.read("a.txt")
    assert initial_text == read_text
    client.mkdir("dir1")
    client.copy("a.txt", "/dir1/b.txt", disk_from="local", disk_to="default")
    read_text = client.read("a.txt", path_to="dir1/b.txt")
    assert "" == read_text
    read_text = client.read("/dir1/b.txt")
    assert read_text == initial_text
    with open(f"{DisksClient.default_disk_root_directory}/dir1/b.txt", "r") as file:
        read_text = file.read()
        assert read_text == initial_text
    os.remove("a.txt")
    client.rm("a.txt")
    client.rm("/dir1", recursive=True)


def test_disks_app_interactive_test_move_and_write():
    initial_text = "File content"
    with open("a.txt", "w") as file:
        file.write(initial_text)
    client = DisksClient.getLocalDisksClient(True)
    client.switch_disk("default")
    client.copy("a.txt", "/a.txt", disk_from="local", disk_to="default")
    files = client.ls(".")
    assert files == ["a.txt"]
    client.move("a.txt", "b.txt")
    files = client.ls(".")
    assert files == ["b.txt"]
    read_text = client.read("/b.txt")
    assert read_text == initial_text
    client.write("b.txt", "c.txt")
    read_text = client.read("c.txt")
    assert read_text == initial_text
    os.remove("a.txt")