def test_request_error_with_prepare_message(self): session = self.make_session() query = SimpleStatement("SELECT * FROM foobar") retry_policy = Mock() retry_policy.on_request_error.return_value = (RetryPolicy.RETHROW, None) message = PrepareMessage(query=query) rf = ResponseFuture(session, message, query, 1, retry_policy=retry_policy) rf._query_retries = 1 rf.send_request() result = Mock(spec=OverloadedErrorMessage) result.to_exception.return_value = result rf._set_result(None, None, None, result) self.assertIsInstance(rf._final_exception, OverloadedErrorMessage) rf = ResponseFuture(session, message, query, 1, retry_policy=retry_policy) rf._query_retries = 1 rf.send_request() result = Mock(spec=ConnectionException) rf._set_result(None, None, None, result) self.assertIsInstance(rf._final_exception, ConnectionException)
def test_errback(self): session = self.make_session() pool = session._pools.get.return_value connection = Mock(spec=Connection) pool.borrow_connection.return_value = (connection, 1) query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf._query_retries = 1 rf.send_request() rf.add_errback(self.assertIsInstance, Exception) result = Mock(spec=UnavailableErrorMessage, info={ "required_replicas": 2, "alive_replicas": 1, "consistency": 1 }) result.to_exception.return_value = Exception() rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result) # this should get called immediately now that the error is set rf.add_errback(self.assertIsInstance, Exception)
def test_unavailable_error_message(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf._query_retries = 1 rf.send_request() result = Mock(spec=UnavailableErrorMessage, info={"required_replicas":2, "alive_replicas": 1, "consistency": 1}) rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result)
def test_add_callbacks(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) # test errback rf = ResponseFuture(session, message, query, 1) rf._query_retries = 1 rf.send_request() rf.add_callbacks(callback=self.assertEqual, callback_args=([{ 'col': 'val' }], ), errback=self.assertIsInstance, errback_args=(Exception, )) result = Mock(spec=UnavailableErrorMessage, info={ "required_replicas": 2, "alive_replicas": 1, "consistency": 1 }) result.to_exception.return_value = Exception() rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result) # test callback rf = ResponseFuture(session, message, query, 1) rf.send_request() callback = Mock() expected_result = (object(), object()) arg = "positional" kwargs = {'one': 1, 'two': 2} rf.add_callbacks(callback=callback, callback_args=(arg, ), callback_kwargs=kwargs, errback=self.assertIsInstance, errback_args=(Exception, )) rf._set_result( None, None, None, self.make_mock_response(expected_result[0], expected_result[1])) self.assertEqual(rf.result()[0], expected_result) callback.assert_called_once_with([expected_result], arg, **kwargs)
def test_unavailable_error_message(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf._query_retries = 1 rf.send_request() result = Mock(spec=UnavailableErrorMessage, info={ "required_replicas": 2, "alive_replicas": 1, "consistency": 1 }) rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result)
def test_errback(self): session = self.make_session() pool = session._pools.get.return_value connection = Mock(spec=Connection) pool.borrow_connection.return_value = (connection, 1) query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf._query_retries = 1 rf.send_request() rf.add_errback(self.assertIsInstance, Exception) result = Mock(spec=UnavailableErrorMessage, info={"required_replicas":2, "alive_replicas": 1, "consistency": 1}) result.to_exception.return_value = Exception() rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result) # this should get called immediately now that the error is set rf.add_errback(self.assertIsInstance, Exception)
def test_add_callbacks(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) # test errback rf = ResponseFuture(session, message, query, 1) rf._query_retries = 1 rf.send_request() rf.add_callbacks( callback=self.assertEqual, callback_args=([{'col': 'val'}],), errback=self.assertIsInstance, errback_args=(Exception,)) result = Mock(spec=UnavailableErrorMessage, info={"required_replicas":2, "alive_replicas": 1, "consistency": 1}) result.to_exception.return_value = Exception() rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result) # test callback rf = ResponseFuture(session, message, query, 1) rf.send_request() callback = Mock() expected_result = [{'col': 'val'}] arg = "positional" kwargs = {'one': 1, 'two': 2} rf.add_callbacks( callback=callback, callback_args=(arg,), callback_kwargs=kwargs, errback=self.assertIsInstance, errback_args=(Exception,)) rf._set_result(None, None, None, self.make_mock_response(expected_result)) self.assertEqual(rf.result(), expected_result) callback.assert_called_once_with(expected_result, arg, **kwargs)