ClickHouse/tests/integration/test_async_drain_connection/test.py

# pylint: disable=redefined-outer-name
# pylint: disable=unused-argument

import pytest
from helpers.cluster import ClickHouseCluster

cluster = ClickHouseCluster(__file__)
node = cluster.add_instance("node", main_configs=["configs/config.xml"])


@pytest.fixture(scope="module")
def started_cluster():
    try:
        cluster.start()
        node.query(
            """
        create table t (number UInt64)
        engine = Distributed(test_cluster_two_shards, system, numbers)
        """
        )
        yield cluster

    finally:
        cluster.shutdown()


def test_filled_async_drain_connection_pool(started_cluster):
    def execute_queries(_):
        for _ in range(100):
            node.query(
                "select * from t where number = 0 limit 2",
                settings={
                    "sleep_in_receive_cancel_ms": int(10e6),
                    "max_execution_time": 5,
                    # decrease drain_timeout to make test more stable
                    # (another way is to increase max_execution_time, but this will make test slower)
                    "drain_timeout": 1,
                },
            )

    any(map(execute_queries, range(10)))