def test_fallback_logic() -> None: state.set_config("use_fallback_host_in_native_connection_pool", 1) network_failure_connection = mock.Mock() network_failure_connection.execute.side_effect = EOFError() verification_connection = mock.Mock() verification_connection.execute.return_value = [] pool = ClickhousePool(CLUSTER_HOST, CLUSTER_PORT, "test", "test", TEST_DB_NAME) # The execute method will try to reuse a single slot in the connection # pool but reestablish new connections with _create_conn if a connection # fails with a network-related error. It may be cleaner to move connection # negotation/establishment into another class for separation of concerns. with mock.patch.object(pool, "_create_conn", lambda x, y=False: network_failure_connection): pool.pool = queue.LifoQueue(1) pool.pool.put(network_failure_connection, block=False) pool.fallback_pool = queue.LifoQueue(1) pool.fallback_pool.put(verification_connection, block=False) pool.execute("SELECT something") assert (network_failure_connection.execute.call_count == 3 ), "Expected three (failed) attempts with main connection pool" assert ( verification_connection.execute.call_count == 1 ), "Expected one (successful) attempt with fallback connection pool"
def test_concurrency_limit() -> None: connection = mock.Mock() connection.execute.side_effect = TestError("some error") state.set_config("simultaneous_queries_sleep_seconds", 0.5) pool = ClickhousePool("host", 100, "test", "test", "test") pool.pool = queue.LifoQueue(1) pool.pool.put(connection, block=False) with pytest.raises(ClickhouseError): pool.execute("SELECT something") connection.execute.assert_called_once() connection.reset_mock(side_effect=True) connection.execute.side_effect = TestConcurrentError("some error") with pytest.raises(ClickhouseError): pool.execute("SELECT something") assert connection.execute.call_count == 2, "Expected two attempts"
def test_robust_concurrency_limit() -> None: connection = mock.Mock() connection.execute.side_effect = ClickhouseError("some error", extra_data={"code": 1}) pool = ClickhousePool("host", 100, "test", "test", "test") pool.pool = queue.LifoQueue(1) pool.pool.put(connection, block=False) with pytest.raises(ClickhouseError): pool.execute_robust("SELECT something") connection.execute.assert_called_once() connection.reset_mock(side_effect=True) connection.execute.side_effect = ClickhouseError( "some error", code=errors.ErrorCodes.TOO_MANY_SIMULTANEOUS_QUERIES, ) with pytest.raises(ClickhouseError): pool.execute_robust("SELECT something") assert connection.execute.call_count == 3, "Expected three attempts"