ClickHouse/tests/integration/test_cgroup_limit/test.py

#!/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()