<commit_before>import Queue
import sys
import threading
import time
import traceback
from collections import OrderedDict
from copy import deepcopy

from cassandra import ConsistencyLevel
from cassandra.query import SimpleStatement

from assertions import assert_none, assert_unavailable
from dtest import DISABLE_VNODES, Tester, debug
from tools import (create_c1c2_table, insert_c1c2, insert_columns, query_c1c2,
                   rows_to_list)


class TestHelper(Tester):

    def __init__(self, *args, **kwargs):
        Tester.__init__(self, *args, **kwargs)
        self.lock = threading.Lock()

    def log(self, message):
        with self.lock:
            debug(message)

    def _name(self, cl):
        return {
            None: '-',
            ConsistencyLevel.ANY: 'ANY',
            ConsistencyLevel.ONE: 'ONE',
            ConsistencyLevel.TWO: 'TWO',
            ConsistencyLevel.THREE: 'THREE',
            ConsistencyLevel.QUORUM: 'QUORUM',
            ConsistencyLevel.ALL: 'ALL',
            ConsistencyLevel.LOCAL_QUORUM: 'LOCAL_QUORUM',
            ConsistencyLevel.EACH_QUORUM: 'EACH_QUORUM',
            ConsistencyLevel.SERIAL: 'SERIAL',
            ConsistencyLevel.LOCAL_SERIAL: 'LOCAL_SERIAL',
            ConsistencyLevel.LOCAL_ONE: 'LOCAL_ONE',
        }[cl]

    def _is_local(self, cl):
        return (cl == ConsistencyLevel.LOCAL_QUORUM or
                cl == ConsistencyLevel.LOCAL_ONE or
                cl == ConsistencyLevel.LOCAL_SERIAL)

    def _is_conditional(self, cl):
        return (cl == ConsistencyLevel.SERIAL or
                cl == ConsistencyLevel.LOCAL_SERIAL)

    def _required_nodes(self, cl, rf_factors, dc):
        """
        Return the number of nodes required by this consistency level
        in the current data center, specified by the dc parameter,
        given a list of replication factors, one per dc.
        """
        return {
            ConsistencyLevel.ANY: 1,
            ConsistencyLevel.ONE: 1,
            ConsistencyLevel.TWO: 2,
            ConsistencyLevel.THREE: 3,
            ConsistencyLevel.QUORUM: sum(rf_factors) / 2 + 1,
            ConsistencyLevel.ALL: sum(rf_factors),
            ConsistencyLevel.LOCAL_QUORUM: rf_factors[dc] / 2 + 1,
            ConsistencyLevel.EACH_QUORUM: rf_factors[dc] / 2 + 1,
            ConsistencyLevel.SERIAL: sum(rf_factors) / 2 + 1,
            ConsistencyLevel.LOCAL_SERIAL: rf_factors[dc] / 2 + 1,
            ConsistencyLevel.LOCAL_ONE: 1,
        }[cl]

    def _should_succeed(self, cl, rf_factors, num_nodes_alive, current):
        """
        Return true if the read or write operation should succeed based on
        the consistency level requested, the replication factors and the
        number of nodes alive in each data center.
        """
        if self._is_local(cl):
            return num_nodes_alive[current] >= self._required_nodes(cl, rf_factors, current)
        elif cl == ConsistencyLevel.EACH_QUORUM:
            for i in xrange(0, len(rf_factors)):
                if num_nodes_alive[i] < self._required_nodes(cl, rf_factors, i):
                    return False
            return True
        else:
            return sum(num_nodes_alive) >= self._required_nodes(cl, rf_factors, current)

    def _start_cluster(self, save_sessions=False):
        cluster = self.cluster
        nodes = self.nodes
        rf = self.rf

        cluster.set_configuration_options(values={'hinted_handoff_enabled': False})
        cluster.populate(nodes).start(wait_for_binary_proto=True, wait_other_notice=True)

        self.ksname = 'mytestks'
        session = self.patient_exclusive_cql_connection(cluster.nodelist()[0])

        self.create_ks(session, self.ksname, rf)
        self.create_tables(session)

        if save_sessions:
            self.sessions = []
            self.sessions.append(session)
            for node in cluster.nodelist()[1:]:
                self.sessions.append(self.patient_exclusive_cql_connection(node, self.ksname))

    def create_tables(self, session):
        self.create_users_table(session)
        self.create_counters_table(session)

    def truncate_tables(self, session):
        statement = SimpleStatement("TRUNCATE users", ConsistencyLevel.ALL)
        session.execute(statement)
        statement = SimpleStatement("TRUNCATE counters", ConsistencyLevel.ALL)
        session.execute(statement)

    def create_users_table(self, session):
        session.execute("""CREATE TABLE users (
                userid int PRIMARY KEY,
                firstname text,
                lastname text,
                age int
            ) WITH COMPACT STORAGE""")

    def insert_user(self, session, userid, age, consistency, serial_consistency=None):
        text = "INSERT INTO users (userid, firstname, lastname, age) VALUES (%d, 'first%d', 'last%d', %d) %s" \
            % (userid, userid, userid, age, "IF NOT EXISTS" if serial_consistency else "")
        statement = SimpleStatement(text, consistency_level=consistency, serial_consistency_level=serial_consistency)
        session.execute(statement)

    def update_user(self, session, userid, age, consistency, serial_consistency=None, prev_age=None):
        text = "UPDATE users SET age = %d WHERE userid = %d" % (age, userid)
        if serial_consistency and prev_age:
            text = text + " IF age = %d" % (prev_age)
        statement = SimpleStatement(text, consistency_level=consistency, serial_consistency_level=serial_consistency)
        session.execute(statement)

    def delete_user(self, session, userid, consistency):
        statement = SimpleStatement("DELETE FROM users where userid = %d" % (userid,), consistency_level=consistency)
        session.execute(statement)

    def query_user(self, session, userid, age, consistency, check_ret=True):
        statement = SimpleStatement("SELECT userid, age FROM users where userid = %d" % (userid,), consistency_level=consistency)
        res = session.execute(statement)
        expected = [[userid, age]] if age else []
        ret = rows_to_list(res) == expected
        if check_ret:
            assert ret, "Got %s from %s, expected %s at %s" % (res, session.cluster.contact_points, expected, self._name(consistency))
        return ret

    def create_counters_table(self, session):
        session.execute("""
            CREATE TABLE counters (
                id int PRIMARY KEY,
                c counter
            )
        """)

    def update_counter(self, session, id, consistency, serial_consistency=None):
        text = "UPDATE counters SET c = c + 1 WHERE id = %d" % (id,)
        statement = SimpleStatement(text, consistency_level=consistency, serial_consistency_level=serial_consistency)<commit_msg>Improve error message.<commit_after>import Queue
import sys
import threading
import time
import traceback
from collections import OrderedDict
from copy import deepcopy

from cassandra import ConsistencyLevel
from cassandra.query import SimpleStatement

from assertions import assert_none, assert_unavailable
from dtest import DISABLE_VNODES, Tester, debug
from tools import (create_c1c2_table, insert_c1c2, insert_columns, query_c1c2,
                   rows_to_list)


class TestHelper(Tester):

    def __init__(self, *args, **kwargs):
        Tester.__init__(self, *args, **kwargs)
        self.lock = threading.Lock()

    def log(self, message):
        with self.lock:
            debug(message)

    def _name(self, cl):
        return {
            None: '-',
            ConsistencyLevel.ANY: 'ANY',
            ConsistencyLevel.ONE: 'ONE',
            ConsistencyLevel.TWO: 'TWO',
            ConsistencyLevel.THREE: 'THREE',
            ConsistencyLevel.QUORUM: 'QUORUM',
            ConsistencyLevel.ALL: 'ALL',
            ConsistencyLevel.LOCAL_QUORUM: 'LOCAL_QUORUM',
            ConsistencyLevel.EACH_QUORUM: 'EACH_QUORUM',
            ConsistencyLevel.SERIAL: 'SERIAL',
            ConsistencyLevel.LOCAL_SERIAL: 'LOCAL_SERIAL',
            ConsistencyLevel.LOCAL_ONE: 'LOCAL_ONE',
        }[cl]

    def _is_local(self, cl):
        return (cl == ConsistencyLevel.LOCAL_QUORUM or
                cl == ConsistencyLevel.LOCAL_ONE or
                cl == ConsistencyLevel.LOCAL_SERIAL)

    def _is_conditional(self, cl):
        return (cl == ConsistencyLevel.SERIAL or
                cl == ConsistencyLevel.LOCAL_SERIAL)

    def _required_nodes(self, cl, rf_factors, dc):
        """
        Return the number of nodes required by this consistency level
        in the current data center, specified by the dc parameter,
        given a list of replication factors, one per dc.
        """
        return {
            ConsistencyLevel.ANY: 1,
            ConsistencyLevel.ONE: 1,
            ConsistencyLevel.TWO: 2,
            ConsistencyLevel.THREE: 3,
            ConsistencyLevel.QUORUM: sum(rf_factors) / 2 + 1,
            ConsistencyLevel.ALL: sum(rf_factors),
            ConsistencyLevel.LOCAL_QUORUM: rf_factors[dc] / 2 + 1,
            ConsistencyLevel.EACH_QUORUM: rf_factors[dc] / 2 + 1,
            ConsistencyLevel.SERIAL: sum(rf_factors) / 2 + 1,
            ConsistencyLevel.LOCAL_SERIAL: rf_factors[dc] / 2 + 1,
            ConsistencyLevel.LOCAL_ONE: 1,
        }[cl]

    def _should_succeed(self, cl, rf_factors, num_nodes_alive, current):
        """
        Return true if the read or write operation should succeed based on
        the consistency level requested, the replication factors and the
        number of nodes alive in each data center.
        """
        if self._is_local(cl):
            return num_nodes_alive[current] >= self._required_nodes(cl, rf_factors, current)
        elif cl == ConsistencyLevel.EACH_QUORUM:
            for i in xrange(0, len(rf_factors)):
                if num_nodes_alive[i] < self._required_nodes(cl, rf_factors, i):
                    return False
            return True
        else:
            return sum(num_nodes_alive) >= self._required_nodes(cl, rf_factors, current)

    def _start_cluster(self, save_sessions=False):
        cluster = self.cluster
        nodes = self.nodes
        rf = self.rf

        cluster.set_configuration_options(values={'hinted_handoff_enabled': False})
        cluster.populate(nodes).start(wait_for_binary_proto=True, wait_other_notice=True)

        self.ksname = 'mytestks'
        session = self.patient_exclusive_cql_connection(cluster.nodelist()[0])

        self.create_ks(session, self.ksname, rf)
        self.create_tables(session)

        if save_sessions:
            self.sessions = []
            self.sessions.append(session)
            for node in cluster.nodelist()[1:]:
                self.sessions.append(self.patient_exclusive_cql_connection(node, self.ksname))

    def create_tables(self, session):
        self.create_users_table(session)
        self.create_counters_table(session)

    def truncate_tables(self, session):
        statement = SimpleStatement("TRUNCATE users", ConsistencyLevel.ALL)
        session.execute(statement)
        statement = SimpleStatement("TRUNCATE counters", ConsistencyLevel.ALL)
        session.execute(statement)

    def create_users_table(self, session):
        session.execute("""CREATE TABLE users (
                userid int PRIMARY KEY,
                firstname text,
                lastname text,
                age int
            ) WITH COMPACT STORAGE""")

    def insert_user(self, session, userid, age, consistency, serial_consistency=None):
        text = "INSERT INTO users (userid, firstname, lastname, age) VALUES (%d, 'first%d', 'last%d', %d) %s" \
            % (userid, userid, userid, age, "IF NOT EXISTS" if serial_consistency else "")
        statement = SimpleStatement(text, consistency_level=consistency, serial_consistency_level=serial_consistency)
        session.execute(statement)

    def update_user(self, session, userid, age, consistency, serial_consistency=None, prev_age=None):
        text = "UPDATE users SET age = %d WHERE userid = %d" % (age, userid)
        if serial_consistency and prev_age:
            text = text + " IF age = %d" % (prev_age)
        statement = SimpleStatement(text, consistency_level=consistency, serial_consistency_level=serial_consistency)
        session.execute(statement)

    def delete_user(self, session, userid, consistency):
        statement = SimpleStatement("DELETE FROM users where userid = %d" % (userid,), consistency_level=consistency)
        session.execute(statement)

    def query_user(self, session, userid, age, consistency, check_ret=True):
        statement = SimpleStatement("SELECT userid, age FROM users where userid = %d" % (userid,), consistency_level=consistency)
        res = session.execute(statement)
        expected = [[userid, age]] if age else []
        ret = rows_to_list(res) == expected
        if check_ret:
            assert ret, "Got %s from %s, expected %s at %s" % (rows_to_list(res), session.cluster.contact_points, expected, self._name(consistency))
        return ret

    def create_counters_table(self, session):
        session.execute("""
            CREATE TABLE counters (
                id int PRIMARY KEY,
                c counter
            )
        """)

    def update_counter(self, session, id, consistency, serial_consistency=None):
        text = "UPDATE counters SET c = c + 1 WHERE id = %d" % (id,)
        statement = SimpleStatement(text, consistency_level=consistency, serial_consistency_level=serial_consistency)
{'fraction_non_alphanumeric': 0.09204183045316325, 'fraction_numerical': 0.003708373507379663, 'mean_word_length': 2.6443243243243244, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6844183', 'n_tokens_mistral': 4063, 'n_tokens_neox': 3856, 'n_words': 957}