Skip to content

Commit bf21403

Browse files
beltranbeltran
committed
Fixed ResponseFuture unit tests
1 parent cc24468 commit bf21403

1 file changed

Lines changed: 28 additions & 16 deletions

File tree

‎tests/unit/test_response_future.py‎

Lines changed: 28 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -39,10 +39,16 @@ class ResponseFutureTests(unittest.TestCase):
3939
def make_basic_session(self):
4040
return Mock(spec=Session, row_factory=lambda *x: list(x))
4141

42+
def make_pool(self):
43+
pool = Mock()
44+
pool.is_shutdown = False
45+
pool.borrow_connection.return_value = [Mock(), Mock()]
46+
return pool
47+
4248
def make_session(self):
4349
session = self.make_basic_session()
4450
session.cluster._default_load_balancing_policy.make_query_plan.return_value = ['ip1', 'ip2']
45-
session._pools.get.return_value.is_shutdown = False
51+
session._pools.get.return_value = self.make_pool()
4652
return session
4753

4854
def make_response_future(self, session):
@@ -122,29 +128,27 @@ def test_other_result_message_kind(self):
122128
def test_read_timeout_error_message(self):
123129
session = self.make_session()
124130
query = SimpleStatement("SELECT * FROM foo")
125-
query.retry_policy = Mock()
126-
query.retry_policy.on_read_timeout.return_value = (RetryPolicy.RETHROW, None)
127131
message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE)
128132

129133
rf = ResponseFuture(session, message, query, 1)
130134
rf.send_request()
131135

132-
result = Mock(spec=ReadTimeoutErrorMessage, info={})
136+
result = Mock(spec=ReadTimeoutErrorMessage, info={"data_retrieved": "", "required_responses":2,
137+
"received_responses":1, "consistency": 1})
133138
rf._set_result(None, None, None, result)
134139

135140
self.assertRaises(Exception, rf.result)
136141

137142
def test_write_timeout_error_message(self):
138143
session = self.make_session()
139144
query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)")
140-
query.retry_policy = Mock()
141-
query.retry_policy.on_write_timeout.return_value = (RetryPolicy.RETHROW, None)
142145
message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE)
143146

144147
rf = ResponseFuture(session, message, query, 1)
145148
rf.send_request()
146149

147-
result = Mock(spec=WriteTimeoutErrorMessage, info={})
150+
result = Mock(spec=WriteTimeoutErrorMessage, info={"write_type": 1, "required_responses":2,
151+
"received_responses":1, "consistency": 1})
148152
rf._set_result(None, None, None, result)
149153
self.assertRaises(Exception, rf.result)
150154

@@ -156,9 +160,10 @@ def test_unavailable_error_message(self):
156160
message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE)
157161

158162
rf = ResponseFuture(session, message, query, 1)
163+
rf._query_retries = 1
159164
rf.send_request()
160165

161-
result = Mock(spec=UnavailableErrorMessage, info={})
166+
result = Mock(spec=UnavailableErrorMessage, info={"required_replicas":2, "alive_replicas": 1, "consistency": 1})
162167
rf._set_result(None, None, None, result)
163168
self.assertRaises(Exception, rf.result)
164169

@@ -295,7 +300,11 @@ def test_first_pool_shutdown(self):
295300
session = self.make_basic_session()
296301
session.cluster._default_load_balancing_policy.make_query_plan.return_value = ['ip1', 'ip2']
297302
# first return a pool with is_shutdown=True, then is_shutdown=False
298-
session._pools.get.side_effect = [Mock(is_shutdown=True), Mock(is_shutdown=False)]
303+
pool_shutdown = self.make_pool()
304+
pool_shutdown.is_shutdown = True
305+
pool_ok = self.make_pool()
306+
pool_ok.is_shutdown = True
307+
session._pools.get.side_effect = [pool_shutdown, pool_ok]
299308

300309
rf = self.make_response_future(session)
301310
rf.send_request()
@@ -358,16 +367,17 @@ def test_errback(self):
358367
pool.borrow_connection.return_value = (connection, 1)
359368

360369
query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)")
361-
query.retry_policy = Mock()
362-
query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None)
363370
message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE)
364371

365372
rf = ResponseFuture(session, message, query, 1)
373+
rf._query_retries = 1
366374
rf.send_request()
367375

368376
rf.add_errback(self.assertIsInstance, Exception)
369377

370-
result = Mock(spec=UnavailableErrorMessage, info={})
378+
result = Mock(spec=UnavailableErrorMessage, info={"required_replicas":2, "alive_replicas": 1, "consistency": 1})
379+
result.to_exception.return_value = Exception()
380+
371381
rf._set_result(None, None, None, result)
372382
self.assertRaises(Exception, rf.result)
373383

@@ -423,9 +433,10 @@ def test_multiple_errbacks(self):
423433
rf.add_errback(callback2, arg2, **kwargs2)
424434

425435
expected_exception = Unavailable("message", 1, 2, 3)
426-
result = Mock(spec=UnavailableErrorMessage, info={'something': 'here'})
436+
result = Mock(spec=UnavailableErrorMessage, info={"required_replicas":2, "alive_replicas": 1, "consistency": 1})
427437
result.to_exception.return_value = expected_exception
428438
rf._set_result(None, None, None, result)
439+
rf._event.set()
429440
self.assertRaises(Exception, rf.result)
430441

431442
callback.assert_called_once_with(expected_exception, arg, **kwargs)
@@ -434,19 +445,20 @@ def test_multiple_errbacks(self):
434445
def test_add_callbacks(self):
435446
session = self.make_session()
436447
query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)")
437-
query.retry_policy = Mock()
438-
query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None)
439448
message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE)
440449

441450
# test errback
442451
rf = ResponseFuture(session, message, query, 1)
452+
rf._query_retries = 1
443453
rf.send_request()
444454

445455
rf.add_callbacks(
446456
callback=self.assertEqual, callback_args=([{'col': 'val'}],),
447457
errback=self.assertIsInstance, errback_args=(Exception,))
448458

449-
result = Mock(spec=UnavailableErrorMessage, info={})
459+
result = Mock(spec=UnavailableErrorMessage,
460+
info={"required_replicas":2, "alive_replicas": 1, "consistency": 1})
461+
result.to_exception.return_value = Exception()
450462
rf._set_result(None, None, None, result)
451463
self.assertRaises(Exception, rf.result)
452464

0 commit comments

Comments
 (0)