1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
|
# Copyright DataStax, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import warnings
import unittest
from cassandra import ConsistencyLevel
from cassandra.cqlengine import connection
from cassandra.cqlengine.management import create_keyspace_simple, drop_keyspace, CQLENG_ALLOW_SCHEMA_MANAGEMENT
import cassandra
from tests.integration import get_server_versions, use_single_node, PROTOCOL_VERSION, CASSANDRA_IP, ALLOW_BETA_PROTOCOL
DEFAULT_KEYSPACE = 'cqlengine_test'
CQL_SKIP_EXECUTE = bool(os.getenv('CQL_SKIP_EXECUTE', False))
def setup_package():
warnings.simplefilter('always') # for testing warnings, make sure all are let through
os.environ[CQLENG_ALLOW_SCHEMA_MANAGEMENT] = '1'
use_single_node()
setup_connection(DEFAULT_KEYSPACE)
create_keyspace_simple(DEFAULT_KEYSPACE, 1)
def teardown_package():
connection.unregister_connection("default")
def is_prepend_reversed():
# do we have https://issues.apache.org/jira/browse/CASSANDRA-8733 ?
ver, _ = get_server_versions()
return not (ver >= (2, 0, 13) or ver >= (2, 1, 3))
def setup_connection(keyspace_name):
connection.setup([CASSANDRA_IP],
consistency=ConsistencyLevel.ONE,
protocol_version=PROTOCOL_VERSION,
allow_beta_protocol_version=ALLOW_BETA_PROTOCOL,
default_keyspace=keyspace_name)
class StatementCounter(object):
"""
Simple python object used to hold a count of the number of times
the wrapped method has been invoked
"""
def __init__(self, patched_func):
self.func = patched_func
self.counter = 0
def wrapped_execute(self, *args, **kwargs):
self.counter += 1
return self.func(*args, **kwargs)
def get_counter(self):
return self.counter
def execute_count(expected):
"""
A decorator used wrap cassandra.cqlengine.connection.execute. It counts the number of times this method is invoked
then compares it to the number expected. If they don't match it throws an assertion error.
This function can be disabled by running the test harness with the env variable CQL_SKIP_EXECUTE=1 set
"""
def innerCounter(fn):
def wrapped_function(*args, **kwargs):
# Create a counter monkey patch into cassandra.cqlengine.connection.execute
count = StatementCounter(cassandra.cqlengine.connection.execute)
original_function = cassandra.cqlengine.connection.execute
# Monkey patch in our StatementCounter wrapper
cassandra.cqlengine.connection.execute = count.wrapped_execute
# Invoked the underlying unit test
to_return = fn(*args, **kwargs)
# Get the count from our monkey patched counter
count.get_counter()
# DeMonkey Patch our code
cassandra.cqlengine.connection.execute = original_function
# Check to see if we have a pre-existing test case to work from.
if len(args) is 0:
test_case = unittest.TestCase("__init__")
else:
test_case = args[0]
# Check to see if the count is what you expect
test_case.assertEqual(count.get_counter(), expected, msg="Expected number of cassandra.cqlengine.connection.execute calls ({0}) doesn't match actual number invoked ({1})".format(expected, count.get_counter()))
return to_return
# Name of the wrapped function must match the original or unittest will error out.
wrapped_function.__name__ = fn.__name__
wrapped_function.__doc__ = fn.__doc__
# Escape hatch
if(CQL_SKIP_EXECUTE):
return fn
else:
return wrapped_function
return innerCounter
|
