Skip to content

Commit 4c0ce05

Browse files
committed
Changing PreparedStatementTests over to use common setup and teardown methods to prevent session leaks
1 parent c91615a commit 4c0ce05

1 file changed

Lines changed: 52 additions & 99 deletions

File tree

tests/integration/standard/test_prepared_statements.py

Lines changed: 52 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -36,21 +36,26 @@ class PreparedStatementTests(unittest.TestCase):
3636
def setUpClass(cls):
3737
cls.cass_version = get_server_versions()
3838

39+
def setUp(self):
40+
self.cluster = Cluster(metrics_enabled=True, protocol_version=PROTOCOL_VERSION)
41+
self.session = self.cluster.connect()
42+
43+
def tearDown(self):
44+
self.cluster.shutdown()
45+
3946
def test_basic(self):
4047
"""
4148
Test basic PreparedStatement usage
4249
"""
4350

44-
cluster = Cluster(protocol_version=PROTOCOL_VERSION)
45-
session = cluster.connect()
46-
session.execute(
51+
self.session.execute(
4752
"""
4853
CREATE KEYSPACE preparedtests
4954
WITH replication = {'class': 'SimpleStrategy', 'replication_factor': '1'}
5055
""")
5156

52-
session.set_keyspace("preparedtests")
53-
session.execute(
57+
self.session.set_keyspace("preparedtests")
58+
self.session.execute(
5459
"""
5560
CREATE TABLE cf0 (
5661
a text,
@@ -60,28 +65,28 @@ def test_basic(self):
6065
)
6166
""")
6267

63-
prepared = session.prepare(
68+
prepared = self.session.prepare(
6469
"""
6570
INSERT INTO cf0 (a, b, c) VALUES (?, ?, ?)
6671
""")
6772

6873
self.assertIsInstance(prepared, PreparedStatement)
6974
bound = prepared.bind(('a', 'b', 'c'))
7075

71-
session.execute(bound)
76+
self.session.execute(bound)
7277

73-
prepared = session.prepare(
78+
prepared = self.session.prepare(
7479
"""
7580
SELECT * FROM cf0 WHERE a=?
7681
""")
7782
self.assertIsInstance(prepared, PreparedStatement)
7883

7984
bound = prepared.bind(('a'))
80-
results = session.execute(bound)
85+
results = self.session.execute(bound)
8186
self.assertEqual(results, [('a', 'b', 'c')])
8287

8388
# test with new dict binding
84-
prepared = session.prepare(
89+
prepared = self.session.prepare(
8590
"""
8691
INSERT INTO cf0 (a, b, c) VALUES (?, ?, ?)
8792
""")
@@ -93,32 +98,26 @@ def test_basic(self):
9398
'c': 'z'
9499
})
95100

96-
session.execute(bound)
101+
self.session.execute(bound)
97102

98-
prepared = session.prepare(
103+
prepared = self.session.prepare(
99104
"""
100105
SELECT * FROM cf0 WHERE a=?
101106
""")
102107

103108
self.assertIsInstance(prepared, PreparedStatement)
104109

105110
bound = prepared.bind({'a': 'x'})
106-
results = session.execute(bound)
111+
results = self.session.execute(bound)
107112
self.assertEqual(results, [('x', 'y', 'z')])
108113

109-
cluster.shutdown()
110-
111114
def test_missing_primary_key(self):
112115
"""
113116
Ensure an InvalidRequest is thrown
114117
when prepared statements are missing the primary key
115118
"""
116119

117-
cluster = Cluster(protocol_version=PROTOCOL_VERSION)
118-
session = cluster.connect()
119-
self._run_missing_primary_key(session)
120-
121-
cluster.shutdown()
120+
self._run_missing_primary_key(self.session)
122121

123122
def _run_missing_primary_key(self, session):
124123
statement_to_prepare = """INSERT INTO test3rf.test (v) VALUES (?)"""
@@ -137,11 +136,7 @@ def test_missing_primary_key_dicts(self):
137136
when prepared statements are missing the primary key
138137
with dict bindings
139138
"""
140-
141-
cluster = Cluster(protocol_version=PROTOCOL_VERSION)
142-
session = cluster.connect()
143-
self._run_missing_primary_key_dicts(session)
144-
cluster.shutdown()
139+
self._run_missing_primary_key_dicts(self.session)
145140

146141
def _run_missing_primary_key_dicts(self, session):
147142
statement_to_prepare = """ INSERT INTO test3rf.test (v) VALUES (?)"""
@@ -158,10 +153,7 @@ def test_too_many_bind_values(self):
158153
"""
159154
Ensure a ValueError is thrown when attempting to bind too many variables
160155
"""
161-
cluster = Cluster(protocol_version=PROTOCOL_VERSION)
162-
session = cluster.connect()
163-
self._run_too_many_bind_values(session)
164-
cluster.shutdown()
156+
self._run_too_many_bind_values(self.session)
165157

166158
def _run_too_many_bind_values(self, session):
167159
statement_to_prepare = """ INSERT INTO test3rf.test (v) VALUES (?)"""
@@ -179,10 +171,7 @@ def test_imprecise_bind_values_dicts(self):
179171
with dict bindings
180172
"""
181173

182-
cluster = Cluster(protocol_version=PROTOCOL_VERSION)
183-
session = cluster.connect()
184-
185-
prepared = session.prepare(
174+
prepared = self.session.prepare(
186175
"""
187176
INSERT INTO test3rf.test (k, v) VALUES (?, ?)
188177
""")
@@ -208,37 +197,30 @@ def test_imprecise_bind_values_dicts(self):
208197
# post v4, the driver attempts to use UNSET_VALUE for unspecified keys
209198
self.assertRaises(ValueError, prepared.bind, {})
210199

211-
cluster.shutdown()
212-
213200
def test_none_values(self):
214201
"""
215202
Ensure binding None is handled correctly
216203
"""
217204

218-
cluster = Cluster(protocol_version=PROTOCOL_VERSION)
219-
session = cluster.connect()
220-
221-
prepared = session.prepare(
205+
prepared = self.session.prepare(
222206
"""
223207
INSERT INTO test3rf.test (k, v) VALUES (?, ?)
224208
""")
225209

226210
self.assertIsInstance(prepared, PreparedStatement)
227211
bound = prepared.bind((1, None))
228-
session.execute(bound)
212+
self.session.execute(bound)
229213

230-
prepared = session.prepare(
214+
prepared = self.session.prepare(
231215
"""
232216
SELECT * FROM test3rf.test WHERE k=?
233217
""")
234218
self.assertIsInstance(prepared, PreparedStatement)
235219

236220
bound = prepared.bind((1,))
237-
results = session.execute(bound)
221+
results = self.session.execute(bound)
238222
self.assertEqual(results[0].v, None)
239223

240-
cluster.shutdown()
241-
242224
def test_unset_values(self):
243225
"""
244226
Test to validate that UNSET_VALUEs are bound, and have the expected effect
@@ -256,13 +238,10 @@ def test_unset_values(self):
256238
if PROTOCOL_VERSION < 4:
257239
raise unittest.SkipTest("Binding UNSET values is not supported in protocol version < 4")
258240

259-
cluster = Cluster(protocol_version=PROTOCOL_VERSION)
260-
session = cluster.connect()
261-
262241
# table with at least two values so one can be used as a marker
263-
session.execute("CREATE TABLE IF NOT EXISTS test1rf.test_unset_values (k int PRIMARY KEY, v0 int, v1 int)")
264-
insert = session.prepare("INSERT INTO test1rf.test_unset_values (k, v0, v1) VALUES (?, ?, ?)")
265-
select = session.prepare("SELECT * FROM test1rf.test_unset_values WHERE k=?")
242+
self.session.execute("CREATE TABLE IF NOT EXISTS test1rf.test_unset_values (k int PRIMARY KEY, v0 int, v1 int)")
243+
insert = self.session.prepare("INSERT INTO test1rf.test_unset_values (k, v0, v1) VALUES (?, ?, ?)")
244+
select = self.session.prepare("SELECT * FROM test1rf.test_unset_values WHERE k=?")
266245

267246
bind_expected = [
268247
# initial condition
@@ -281,127 +260,105 @@ def test_unset_values(self):
281260
]
282261

283262
for params, expected in bind_expected:
284-
session.execute(insert, params)
285-
results = session.execute(select, (0,))
263+
self.session.execute(insert, params)
264+
results = self.session.execute(select, (0,))
286265
self.assertEqual(results[0], expected)
287266

288-
self.assertRaises(ValueError, session.execute, select, (UNSET_VALUE, 0, 0))
289-
290-
cluster.shutdown()
267+
self.assertRaises(ValueError, self.session.execute, select, (UNSET_VALUE, 0, 0))
291268

292269
def test_no_meta(self):
293-
cluster = Cluster(protocol_version=PROTOCOL_VERSION)
294-
session = cluster.connect()
295270

296-
prepared = session.prepare(
271+
prepared = self.session.prepare(
297272
"""
298273
INSERT INTO test3rf.test (k, v) VALUES (0, 0)
299274
""")
300275

301276
self.assertIsInstance(prepared, PreparedStatement)
302277
bound = prepared.bind(None)
303-
session.execute(bound)
278+
self.session.execute(bound)
304279

305-
prepared = session.prepare(
280+
prepared = self.session.prepare(
306281
"""
307282
SELECT * FROM test3rf.test WHERE k=0
308283
""")
309284
self.assertIsInstance(prepared, PreparedStatement)
310285

311286
bound = prepared.bind(None)
312-
results = session.execute(bound)
287+
results = self.session.execute(bound)
313288
self.assertEqual(results[0].v, 0)
314289

315-
cluster.shutdown()
316-
317290
def test_none_values_dicts(self):
318291
"""
319292
Ensure binding None is handled correctly with dict bindings
320293
"""
321294

322-
cluster = Cluster(protocol_version=PROTOCOL_VERSION)
323-
session = cluster.connect()
324-
325295
# test with new dict binding
326-
prepared = session.prepare(
296+
prepared = self.session.prepare(
327297
"""
328298
INSERT INTO test3rf.test (k, v) VALUES (?, ?)
329299
""")
330300

331301
self.assertIsInstance(prepared, PreparedStatement)
332302
bound = prepared.bind({'k': 1, 'v': None})
333-
session.execute(bound)
303+
self.session.execute(bound)
334304

335-
prepared = session.prepare(
305+
prepared = self.session.prepare(
336306
"""
337307
SELECT * FROM test3rf.test WHERE k=?
338308
""")
339309
self.assertIsInstance(prepared, PreparedStatement)
340310

341311
bound = prepared.bind({'k': 1})
342-
results = session.execute(bound)
312+
results = self.session.execute(bound)
343313
self.assertEqual(results[0].v, None)
344314

345-
cluster.shutdown()
346-
347315
def test_async_binding(self):
348316
"""
349317
Ensure None binding over async queries
350318
"""
351319

352-
cluster = Cluster(protocol_version=PROTOCOL_VERSION)
353-
session = cluster.connect()
354-
355-
prepared = session.prepare(
320+
prepared = self.session.prepare(
356321
"""
357322
INSERT INTO test3rf.test (k, v) VALUES (?, ?)
358323
""")
359324

360325
self.assertIsInstance(prepared, PreparedStatement)
361-
future = session.execute_async(prepared, (873, None))
326+
future = self.session.execute_async(prepared, (873, None))
362327
future.result()
363328

364-
prepared = session.prepare(
329+
prepared = self.session.prepare(
365330
"""
366331
SELECT * FROM test3rf.test WHERE k=?
367332
""")
368333
self.assertIsInstance(prepared, PreparedStatement)
369334

370-
future = session.execute_async(prepared, (873,))
335+
future = self.session.execute_async(prepared, (873,))
371336
results = future.result()
372337
self.assertEqual(results[0].v, None)
373338

374-
cluster.shutdown()
375-
376339
def test_async_binding_dicts(self):
377340
"""
378341
Ensure None binding over async queries with dict bindings
379342
"""
380-
381-
cluster = Cluster(protocol_version=PROTOCOL_VERSION)
382-
session = cluster.connect()
383-
384-
prepared = session.prepare(
343+
prepared = self.session.prepare(
385344
"""
386345
INSERT INTO test3rf.test (k, v) VALUES (?, ?)
387346
""")
388347

389348
self.assertIsInstance(prepared, PreparedStatement)
390-
future = session.execute_async(prepared, {'k': 873, 'v': None})
349+
future = self.session.execute_async(prepared, {'k': 873, 'v': None})
391350
future.result()
392351

393-
prepared = session.prepare(
352+
prepared = self.session.prepare(
394353
"""
395354
SELECT * FROM test3rf.test WHERE k=?
396355
""")
397356
self.assertIsInstance(prepared, PreparedStatement)
398357

399-
future = session.execute_async(prepared, {'k': 873})
358+
future = self.session.execute_async(prepared, {'k': 873})
400359
results = future.result()
401360
self.assertEqual(results[0].v, None)
402361

403-
cluster.shutdown()
404-
405362
def test_raise_error_on_prepared_statement_execution_dropped_table(self):
406363
"""
407364
test for error in executing prepared statement on a dropped table
@@ -418,14 +375,10 @@ def test_raise_error_on_prepared_statement_execution_dropped_table(self):
418375
419376
@test_category prepared_statements
420377
"""
421-
cluster = Cluster(protocol_version=PROTOCOL_VERSION)
422-
session = cluster.connect("test3rf")
423378

424-
session.execute("CREATE TABLE error_test (k int PRIMARY KEY, v int)")
425-
prepared = session.prepare("SELECT * FROM error_test WHERE k=?")
426-
session.execute("DROP TABLE error_test")
379+
self.session.execute("CREATE TABLE test3rf.error_test (k int PRIMARY KEY, v int)")
380+
prepared = self.session.prepare("SELECT * FROM test3rf.error_test WHERE k=?")
381+
self.session.execute("DROP TABLE test3rf.error_test")
427382

428383
with self.assertRaises(InvalidRequest):
429-
session.execute(prepared, [0])
430-
431-
cluster.shutdown()
384+
self.session.execute(prepared, [0])

0 commit comments

Comments
 (0)