#!/usr/bin/env python3

import os
import math
import subprocess
from tempfile import NamedTemporaryFile
import pytest


def run_command_in_container(cmd, *args):
    # /clickhouse is mounted by integration tests runner
    alternative_binary = os.getenv("CLICKHOUSE_BINARY", "/clickhouse")
    if alternative_binary:
        args += (
            "--volume",
            f"{alternative_binary}:/usr/bin/clickhouse",
        )

    return subprocess.check_output(
        [
            "docker",
            "run",
            "--rm",
            *args,
            "ubuntu:20.04",
            "sh",
            "-c",
            cmd,
        ]
    )


def run_with_cpu_limit(cmd, num_cpus, *args):
    args += (
        "--cpus",
        f"{num_cpus}",
    )
    return run_command_in_container(cmd, *args)


def test_cgroup_cpu_limit():
    for num_cpus in (1, 2, 4, 2.8):
        result = run_with_cpu_limit(
            "clickhouse local -q \"select value from system.settings where name='max_threads'\"",
            num_cpus,
        )
        expect_output = (r"\'auto({})\'".format(math.ceil(num_cpus))).encode()
        assert (
            result.strip() == expect_output
        ), f"fail for cpu limit={num_cpus}, result={result.strip()}, expect={expect_output}"


# For manual run
if __name__ == "__main__":
    test_cgroup_cpu_limit()