Skip to content

Commit 47cb3e9

Browse files
GregBestlandGregBestland
committed
Adding decorator for counting underlying query execution, updating some tests to add expected query counts
1 parent c7f15a2 commit 47cb3e9

7 files changed

Lines changed: 162 additions & 13 deletions

File tree

‎tests/integration/cqlengine/__init__.py‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,23 @@
1414

1515
import os
1616
import warnings
17+
try:
18+
import unittest2 as unittest
19+
except ImportError:
20+
import unittest # noqa
1721
from cassandra import ConsistencyLevel
1822

1923
from cassandra.cqlengine import connection
2024
from cassandra.cqlengine.management import create_keyspace_simple, CQLENG_ALLOW_SCHEMA_MANAGEMENT
25+
import cassandra
2126

2227
from tests.integration import get_server_versions, use_single_node, PROTOCOL_VERSION
2328
DEFAULT_KEYSPACE = 'cqlengine_test'
2429

2530

31+
CQL_SKIP_EXECUTE = bool(os.getenv('CQL_SKIP_EXECUTE', False))
32+
33+
2634
def setup_package():
2735
warnings.simplefilter('always') # for testing warnings, make sure all are let through
2836
os.environ[CQLENG_ALLOW_SCHEMA_MANAGEMENT] = '1'
@@ -44,3 +52,57 @@ def setup_connection(keyspace_name):
4452
consistency=ConsistencyLevel.ONE,
4553
protocol_version=PROTOCOL_VERSION,
4654
default_keyspace=keyspace_name)
55+
56+
57+
class StatementCounter(object):
58+
"""
59+
Simple python object used to hold a count of the number of times
60+
the wrapped method has been invoked
61+
"""
62+
def __init__(self, patched_func):
63+
self.func = patched_func
64+
self.counter = 0
65+
66+
def wrapped_execute(self, *args, **kwargs):
67+
self.counter += 1
68+
return self.func(*args, **kwargs)
69+
70+
def get_counter(self):
71+
return self.counter
72+
73+
74+
def execute_count(expected):
75+
"""
76+
A decorator used wrap cassandra.cqlengine.connection.execute. It counts the number of times this method is invoked
77+
then compares it to the number expected. If they don't match it throws an assertion error.
78+
This function can be disabled by running the test harness with the env variable CQL_SKIP_EXECUTE=1 set
79+
"""
80+
def innerCounter(fn):
81+
def wrapped_function(*args, **kwargs):
82+
# Create a counter monkey patch into cassandra.cqlengine.connection.execute
83+
count = StatementCounter(cassandra.cqlengine.connection.execute)
84+
original_function = cassandra.cqlengine.connection.execute
85+
# Monkey patch in our StatementCounter wrapper
86+
cassandra.cqlengine.connection.execute = count.wrapped_execute
87+
# Invoked the underlying unit test
88+
to_return = fn(*args, **kwargs)
89+
# Get the count from our monkey patched counter
90+
count.get_counter()
91+
# DeMonkey Patch our code
92+
cassandra.cqlengine.connection.execute = original_function
93+
# Check to see if the count is what you expect
94+
tc = unittest.TestCase("__init__")
95+
tc.assertEquals(count.get_counter(), expected, "Expected number of cassandra.cqlengine.connection.execute calls doesn't match actual number invoked Expected: {0}, Invoked {1}".format(count.get_counter(), expected))
96+
return to_return
97+
# Name of the wrapped function must match the original or unittest will error out.
98+
wrapped_function.__name__ = fn.__name__
99+
wrapped_function.__doc__ = fn.__doc__
100+
# Escape hatch
101+
if(CQL_SKIP_EXECUTE):
102+
return fn
103+
else:
104+
return wrapped_function
105+
106+
return innerCounter
107+
108+

‎tests/integration/cqlengine/query/test_batch_query.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from cassandra.cqlengine.models import Model
2121
from cassandra.cqlengine.query import BatchQuery, DMLQuery
2222
from tests.integration.cqlengine.base import BaseCassEngTestCase
23+
from tests.integration.cqlengine import execute_count
2324
from cassandra.cluster import Session
2425

2526

@@ -55,6 +56,7 @@ def setUp(self):
5556
for obj in TestMultiKeyModel.filter(partition=self.pkey):
5657
obj.delete()
5758

59+
@execute_count(3)
5860
def test_insert_success_case(self):
5961

6062
b = BatchQuery()
@@ -67,6 +69,7 @@ def test_insert_success_case(self):
6769

6870
TestMultiKeyModel.get(partition=self.pkey, cluster=2)
6971

72+
@execute_count(4)
7073
def test_update_success_case(self):
7174

7275
inst = TestMultiKeyModel.create(partition=self.pkey, cluster=2, count=3, text='4')
@@ -84,6 +87,7 @@ def test_update_success_case(self):
8487
inst3 = TestMultiKeyModel.get(partition=self.pkey, cluster=2)
8588
self.assertEqual(inst3.count, 4)
8689

90+
@execute_count(4)
8791
def test_delete_success_case(self):
8892

8993
inst = TestMultiKeyModel.create(partition=self.pkey, cluster=2, count=3, text='4')
@@ -99,6 +103,7 @@ def test_delete_success_case(self):
99103
with self.assertRaises(TestMultiKeyModel.DoesNotExist):
100104
TestMultiKeyModel.get(partition=self.pkey, cluster=2)
101105

106+
@execute_count(11)
102107
def test_context_manager(self):
103108

104109
with BatchQuery() as b:
@@ -112,6 +117,7 @@ def test_context_manager(self):
112117
for i in range(5):
113118
TestMultiKeyModel.get(partition=self.pkey, cluster=i)
114119

120+
@execute_count(9)
115121
def test_bulk_delete_success_case(self):
116122

117123
for i in range(1):
@@ -127,6 +133,7 @@ def test_bulk_delete_success_case(self):
127133
for m in TestMultiKeyModel.all():
128134
m.delete()
129135

136+
@execute_count(0)
130137
def test_none_success_case(self):
131138
""" Tests that passing None into the batch call clears any batch object """
132139
b = BatchQuery()
@@ -137,6 +144,7 @@ def test_none_success_case(self):
137144
q = q.batch(None)
138145
self.assertIsNone(q._batch)
139146

147+
@execute_count(0)
140148
def test_dml_none_success_case(self):
141149
""" Tests that passing None into the batch call clears any batch object """
142150
b = BatchQuery()
@@ -147,6 +155,7 @@ def test_dml_none_success_case(self):
147155
q.batch(None)
148156
self.assertIsNone(q._batch)
149157

158+
@execute_count(3)
150159
def test_batch_execute_on_exception_succeeds(self):
151160
# makes sure if execute_on_exception == True we still apply the batch
152161
drop_table(BatchQueryLogModel)
@@ -166,6 +175,7 @@ def test_batch_execute_on_exception_succeeds(self):
166175
# should be 1 because the batch should execute
167176
self.assertEqual(1, len(obj))
168177

178+
@execute_count(2)
169179
def test_batch_execute_on_exception_skips_if_not_specified(self):
170180
# makes sure if execute_on_exception == True we still apply the batch
171181
drop_table(BatchQueryLogModel)
@@ -186,12 +196,14 @@ def test_batch_execute_on_exception_skips_if_not_specified(self):
186196
# should be 0 because the batch should not execute
187197
self.assertEqual(0, len(obj))
188198

199+
@execute_count(1)
189200
def test_batch_execute_timeout(self):
190201
with mock.patch.object(Session, 'execute') as mock_execute:
191202
with BatchQuery(timeout=1) as b:
192203
BatchQueryLogModel.batch(b).create(k=2, v=2)
193204
self.assertEqual(mock_execute.call_args[-1]['timeout'], 1)
194205

206+
@execute_count(1)
195207
def test_batch_execute_no_timeout(self):
196208
with mock.patch.object(Session, 'execute') as mock_execute:
197209
with BatchQuery() as b:

‎tests/integration/cqlengine/query/test_datetime_queries.py‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,15 +20,17 @@
2020

2121
from cassandra.cqlengine.management import sync_table
2222
from cassandra.cqlengine.management import drop_table
23-
from cassandra.cqlengine.models import Model, ModelException
23+
from cassandra.cqlengine.models import Model
2424
from cassandra.cqlengine import columns
25-
from cassandra.cqlengine import query
25+
from tests.integration.cqlengine import execute_count
26+
2627

2728
class DateTimeQueryTestModel(Model):
2829

29-
user = columns.Integer(primary_key=True)
30-
day = columns.DateTime(primary_key=True)
31-
data = columns.Text()
30+
user = columns.Integer(primary_key=True)
31+
day = columns.DateTime(primary_key=True)
32+
data = columns.Text()
33+
3234

3335
class TestDateTimeQueries(BaseCassEngTestCase):
3436

@@ -46,12 +48,12 @@ def setUpClass(cls):
4648
data=str(uuid4())
4749
)
4850

49-
5051
@classmethod
5152
def tearDownClass(cls):
5253
super(TestDateTimeQueries, cls).tearDownClass()
5354
drop_table(DateTimeQueryTestModel)
5455

56+
@execute_count(1)
5557
def test_range_query(self):
5658
""" Tests that loading from a range of dates works properly """
5759
start = datetime(*self.base_date.timetuple()[:3])
@@ -60,6 +62,7 @@ def test_range_query(self):
6062
results = DateTimeQueryTestModel.filter(user=0, day__gte=start, day__lt=end)
6163
assert len(results) == 3
6264

65+
@execute_count(3)
6366
def test_datetime_precision(self):
6467
""" Tests that millisecond resolution is preserved when saving datetime objects """
6568
now = datetime.now()

‎tests/integration/cqlengine/query/test_named.py‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
from cassandra.concurrent import execute_concurrent_with_args
2626
from cassandra.cqlengine import models
2727

28-
from tests.integration.cqlengine import setup_connection
28+
from tests.integration.cqlengine import setup_connection, execute_count
2929
from tests.integration.cqlengine.base import BaseCassEngTestCase
3030
from tests.integration.cqlengine.query.test_queryset import BaseQuerySetUsage
3131

@@ -134,20 +134,23 @@ def setUpClass(cls):
134134
cls.keyspace = NamedKeyspace(ks)
135135
cls.table = cls.keyspace.table(tn)
136136

137+
@execute_count(2)
137138
def test_count(self):
138139
""" Tests that adding filtering statements affects the count query as expected """
139140
assert self.table.objects.count() == 12
140141

141142
q = self.table.objects(test_id=0)
142143
assert q.count() == 4
143144

145+
@execute_count(2)
144146
def test_query_expression_count(self):
145147
""" Tests that adding query statements affects the count query as expected """
146148
assert self.table.objects.count() == 12
147149

148150
q = self.table.objects(self.table.column('test_id') == 0)
149151
assert q.count() == 4
150152

153+
@execute_count(3)
151154
def test_iteration(self):
152155
""" Tests that iterating over a query set pulls back all of the expected results """
153156
q = self.table.objects(test_id=0)
@@ -181,6 +184,7 @@ def test_iteration(self):
181184
compare_set.remove(val)
182185
assert len(compare_set) == 0
183186

187+
@execute_count(2)
184188
def test_multiple_iterations_work_properly(self):
185189
""" Tests that iterating over a query set more than once works """
186190
# test with both the filtering method and the query method
@@ -201,6 +205,7 @@ def test_multiple_iterations_work_properly(self):
201205
compare_set.remove(val)
202206
assert len(compare_set) == 0
203207

208+
@execute_count(2)
204209
def test_multiple_iterators_are_isolated(self):
205210
"""
206211
tests that the use of one iterator does not affect the behavior of another
@@ -214,6 +219,7 @@ def test_multiple_iterators_are_isolated(self):
214219
assert next(iter1).attempt_id == attempt_id
215220
assert next(iter2).attempt_id == attempt_id
216221

222+
@execute_count(3)
217223
def test_get_success_case(self):
218224
"""
219225
Tests that the .get() method works on new and existing querysets
@@ -235,6 +241,7 @@ def test_get_success_case(self):
235241
assert m.test_id == 0
236242
assert m.attempt_id == 0
237243

244+
@execute_count(3)
238245
def test_query_expression_get_success_case(self):
239246
"""
240247
Tests that the .get() method works on new and existing querysets
@@ -256,13 +263,15 @@ def test_query_expression_get_success_case(self):
256263
assert m.test_id == 0
257264
assert m.attempt_id == 0
258265

266+
@execute_count(1)
259267
def test_get_doesnotexist_exception(self):
260268
"""
261269
Tests that get calls that don't return a result raises a DoesNotExist error
262270
"""
263271
with self.assertRaises(self.table.DoesNotExist):
264272
self.table.objects.get(test_id=100)
265273

274+
@execute_count(1)
266275
def test_get_multipleobjects_exception(self):
267276
"""
268277
Tests that get calls that return multiple results raise a MultipleObjectsReturned error
@@ -286,6 +295,7 @@ def tearDownClass(cls):
286295
super(TestNamedWithMV, cls).tearDownClass()
287296

288297
@greaterthanorequalcass30
298+
@execute_count(5)
289299
def test_named_table_with_mv(self):
290300
"""
291301
Test NamedTable access to materialized views

‎tests/integration/cqlengine/query/test_queryoperators.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from cassandra.cqlengine.statements import WhereClause
2525
from tests.integration.cqlengine import DEFAULT_KEYSPACE
2626
from tests.integration.cqlengine.base import BaseCassEngTestCase
27+
from tests.integration.cqlengine import execute_count
2728

2829

2930
class TestQuerySetOperation(BaseCassEngTestCase):
@@ -71,6 +72,7 @@ def tearDown(self):
7172
super(TestTokenFunction, self).tearDown()
7273
drop_table(TokenTestModel)
7374

75+
@execute_count(14)
7476
def test_token_function(self):
7577
""" Tests that token functions work properly """
7678
assert TokenTestModel.objects().count() == 0
@@ -127,6 +129,7 @@ class TestModel(Model):
127129
func = functions.Token('a')
128130
self.assertRaises(query.QueryException, TestModel.objects.filter, pk__token__gt=func)
129131

132+
@execute_count(7)
130133
def test_named_table_pk_token_function(self):
131134
"""
132135
Test to ensure that token function work with named tables.

0 commit comments

Comments
 (0)