54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
import logging
|
|
log = logging.getLogger(__name__)
|
|
import unittest
|
|
|
|
from cassandra.cluster import Cluster
|
|
from cassandra.policies import HostDistance
|
|
|
|
existing_keyspaces = None
|
|
|
|
def setup_package():
|
|
try:
|
|
cluster = Cluster()
|
|
cluster.set_core_connections_per_host(HostDistance.LOCAL, 1)
|
|
cluster.set_max_connections_per_host(HostDistance.LOCAL, 1)
|
|
session = cluster.connect()
|
|
except Exception, exc:
|
|
log.exception('Failed to connect to cluster:')
|
|
raise unittest.SkipTest('Failed to connect to cluster: %r' % exc)
|
|
|
|
try:
|
|
global existing_keyspaces
|
|
results = session.execute("SELECT keyspace_name FROM system.schema_keyspaces")
|
|
existing_keyspaces = set([row[0] for row in results])
|
|
finally:
|
|
try:
|
|
cluster.shutdown()
|
|
except Exception, exc:
|
|
log.exception('Failed to connect to cluster:')
|
|
raise unittest.SkipTest('Failed to connect to cluster: %r' % exc)
|
|
|
|
|
|
def teardown_package():
|
|
try:
|
|
cluster = Cluster()
|
|
cluster.set_core_connections_per_host(HostDistance.LOCAL, 1)
|
|
cluster.set_max_connections_per_host(HostDistance.LOCAL, 1)
|
|
session = cluster.connect()
|
|
except Exception:
|
|
log.exception('Failed to connect to cluster:')
|
|
return
|
|
|
|
try:
|
|
if existing_keyspaces:
|
|
results = session.execute("SELECT keyspace_name FROM system.schema_keyspaces")
|
|
current_keyspaces = set([row[0] for row in results])
|
|
for keyspace in current_keyspaces - existing_keyspaces:
|
|
session.execute("DROP KEYSPACE %s" % (keyspace,))
|
|
|
|
finally:
|
|
try:
|
|
cluster.shutdown()
|
|
except:
|
|
log.exception('Failed to connect to cluster:')
|