<commit_before>import struct
import time

from cassandra.query import SimpleStatement
from cassandra import ConsistencyLevel
from tests.integration import get_node


coordinators = {}


def add_coordinator(future):
    global coordinators
    coordinator = future._current_host.address
    if coordinator in coordinators:
        coordinators[coordinator] += 1
    else:
        coordinators[coordinator] = 1
    if future._errors:
        print 'future._errors', future._errors
    future.result()


def reset_coordinators():
    global coordinators
    coordinators = {}


def get_queried(node):
    ip = '127.0.0.%s' % node
    if not ip in coordinators:
        return 0
    return coordinators[ip]


def assert_queried(node, n):
    ip = '127.0.0.%s' % node
    if ip in coordinators:
        if coordinators[ip] == n:
            return
        raise RuntimeError(
            'IP: %s. Expected: %s. Received: %s. Full detail: %s.' % (ip, n, coordinators[ip], coordinators))
    else:
        if n == 0:
            return
        raise RuntimeError('IP: %s. Expected: %s. Received: %s. Full detail: %s.' % (ip, n, 0, coordinators))


def create_schema(session, keyspace, replication_class='SimpleStrategy',
                  replication_factor=1):

    results = session.execute(
        'SELECT keyspace_name FROM system.schema_keyspaces')
    existing_keyspaces = [row[0] for row in results]
    if keyspace in existing_keyspaces:
        session.execute('DROP KEYSPACE %s' % keyspace)

    if replication_class == 'SimpleStrategy':
        ddl = "CREATE KEYSPACE %s WITH replication" \
              " = {'class': 'SimpleStrategy', 'replication_factor': '%s'}"
        session.execute(ddl % (keyspace, replication_factor))

        ddl = 'CREATE TABLE %s.cf (k int PRIMARY KEY, i int)'
        session.execute(ddl % keyspace)
    session.execute('USE %s' % keyspace)

    # BUG: probably related to PYTHON-39
    time.sleep(5)


def init(session, keyspace, n, consistency_level=ConsistencyLevel.ONE):
    reset_coordinators()
    # BUG: PYTHON-38
    # session.execute('USE %s' % keyspace)
    for i in range(n):
        ss = SimpleStatement('INSERT INTO %s(k, i) VALUES (0, 0)' % 'cf',
                             consistency_level=consistency_level)
        session.execute(ss)


def query(session, keyspace, n, consistency_level=ConsistencyLevel.ONE):
    routing_key = struct.pack('>i', 0)
    for i in range(n):
        ss = SimpleStatement('SELECT * FROM %s WHERE k = 0' % 'cf',
                             consistency_level=consistency_level,
                             routing_key=routing_key)
        add_coordinator(session.execute_async(ss))
<commit_msg>Modify create_schema to work with NTS.<commit_after>import struct
import time

from cassandra.query import SimpleStatement
from cassandra import ConsistencyLevel
from tests.integration import get_node


coordinators = {}


def add_coordinator(future):
    global coordinators
    coordinator = future._current_host.address
    if coordinator in coordinators:
        coordinators[coordinator] += 1
    else:
        coordinators[coordinator] = 1
    if future._errors:
        print 'future._errors', future._errors
    future.result()


def reset_coordinators():
    global coordinators
    coordinators = {}


def get_queried(node):
    ip = '127.0.0.%s' % node
    if not ip in coordinators:
        return 0
    return coordinators[ip]


def assert_queried(node, n):
    ip = '127.0.0.%s' % node
    if ip in coordinators:
        if coordinators[ip] == n:
            return
        raise RuntimeError(
            'IP: %s. Expected: %s. Received: %s. Full detail: %s.' % (ip, n, coordinators[ip], coordinators))
    else:
        if n == 0:
            return
        raise RuntimeError('IP: %s. Expected: %s. Received: %s. Full detail: %s.' % (ip, n, 0, coordinators))


def create_schema(session, keyspace, replication_class='SS',
                  replication_factor=1, replication_strategy=None):

    results = session.execute(
        'SELECT keyspace_name FROM system.schema_keyspaces')
    existing_keyspaces = [row[0] for row in results]
    if keyspace in existing_keyspaces:
        session.execute('DROP KEYSPACE %s' % keyspace)

    if replication_class == 'SS':
        ddl = "CREATE KEYSPACE %s WITH replication" \
              " = {'class': 'SimpleStrategy', 'replication_factor': '%s'}"
        session.execute(ddl % (keyspace, replication_factor))
    elif replication_class == 'NTS':
        if not replication_strategy:
            raise Exception('replication_strategy is not set')

        ddl = "CREATE KEYSPACE %s" \
              " WITH replication = { 'class' : 'NetworkTopologyStrategy', %s }"
        session.execute(ddl % (keyspace, str(replication_strategy)[1:-1]))

    ddl = 'CREATE TABLE %s.cf (k int PRIMARY KEY, i int)'
    session.execute(ddl % keyspace)
    session.execute('USE %s' % keyspace)

    # BUG: probably related to PYTHON-39
    time.sleep(5)


def init(session, keyspace, n, consistency_level=ConsistencyLevel.ONE):
    reset_coordinators()
    # BUG: PYTHON-38
    # session.execute('USE %s' % keyspace)
    for i in range(n):
        ss = SimpleStatement('INSERT INTO %s(k, i) VALUES (0, 0)' % 'cf',
                             consistency_level=consistency_level)
        session.execute(ss)


def query(session, keyspace, n, consistency_level=ConsistencyLevel.ONE):
    routing_key = struct.pack('>i', 0)
    for i in range(n):
        ss = SimpleStatement('SELECT * FROM %s WHERE k = 0' % 'cf',
                             consistency_level=consistency_level,
                             routing_key=routing_key)
        add_coordinator(session.execute_async(ss))

{'fraction_non_alphanumeric': 0.11019621583742116, 'fraction_numerical': 0.0094604064470918, 'mean_word_length': 2.726501305483029, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '4099919', 'n_tokens_mistral': 1823, 'n_tokens_neox': 1744, 'n_words': 461}