227 lines
7.6 KiB
Python
227 lines
7.6 KiB
Python
import errno
|
|
from StringIO import StringIO
|
|
from socket import error as socket_error
|
|
import unittest
|
|
|
|
from mock import patch, Mock
|
|
|
|
from cassandra.connection import (PROTOCOL_VERSION,
|
|
HEADER_DIRECTION_TO_CLIENT,
|
|
ProtocolError,
|
|
ConnectionException)
|
|
from cassandra.io.libevreactor import LibevConnection
|
|
from cassandra.decoder import (write_stringmultimap, write_int, write_string,
|
|
SupportedMessage, ReadyMessage, ServerError)
|
|
from cassandra.marshal import uint8_pack, uint32_pack
|
|
|
|
@patch('socket.socket')
|
|
@patch('pyev.Io')
|
|
@patch('cassandra.io.asyncorereactor._start_loop')
|
|
class LibevConnectionTest(unittest.TestCase):
|
|
|
|
def make_connection(self):
|
|
c = LibevConnection('1.2.3.4')
|
|
c._socket = Mock()
|
|
c._socket.send.side_effect = lambda x: len(x)
|
|
return c
|
|
|
|
def make_header_prefix(self, message_class, version=PROTOCOL_VERSION, stream_id=0):
|
|
return ''.join(map(uint8_pack, [
|
|
0xff & (HEADER_DIRECTION_TO_CLIENT | version),
|
|
0, # flags (compression)
|
|
stream_id,
|
|
message_class.opcode # opcode
|
|
]))
|
|
|
|
def make_options_body(self):
|
|
options_buf = StringIO()
|
|
write_stringmultimap(options_buf, {
|
|
'CQL_VERSION': ['3.0.1'],
|
|
'COMPRESSION': []
|
|
})
|
|
return options_buf.getvalue()
|
|
|
|
def make_error_body(self, code, msg):
|
|
buf = StringIO()
|
|
write_int(buf, code)
|
|
write_string(buf, msg)
|
|
return buf.getvalue()
|
|
|
|
def make_msg(self, header, body=""):
|
|
return header + uint32_pack(len(body)) + body
|
|
|
|
def test_successful_connection(self, *args):
|
|
c = self.make_connection()
|
|
|
|
# let it write the OptionsMessage
|
|
c.handle_write(None, None)
|
|
|
|
# read in a SupportedMessage response
|
|
header = self.make_header_prefix(SupportedMessage)
|
|
options = self.make_options_body()
|
|
c._socket.recv.return_value = self.make_msg(header, options)
|
|
c.handle_read(None, None)
|
|
|
|
# let it write out a StartupMessage
|
|
c.handle_write(None, None)
|
|
|
|
header = self.make_header_prefix(ReadyMessage, stream_id=1)
|
|
c._socket.recv.return_value = self.make_msg(header)
|
|
c.handle_read(None, None)
|
|
|
|
self.assertTrue(c.connected_event.is_set())
|
|
|
|
def test_protocol_error(self, *args):
|
|
c = self.make_connection()
|
|
|
|
# let it write the OptionsMessage
|
|
c.handle_write(None, None)
|
|
|
|
# read in a SupportedMessage response
|
|
header = self.make_header_prefix(SupportedMessage, version=0x04)
|
|
options = self.make_options_body()
|
|
c._socket.recv.return_value = self.make_msg(header, options)
|
|
c.handle_read(None, None)
|
|
|
|
# make sure it errored correctly
|
|
self.assertTrue(c.is_defunct)
|
|
self.assertTrue(isinstance(c.last_error, ProtocolError))
|
|
self.assertTrue(c.connected_event.is_set())
|
|
|
|
def test_error_message_on_startup(self, *args):
|
|
c = self.make_connection()
|
|
|
|
# let it write the OptionsMessage
|
|
c.handle_write(None, None)
|
|
|
|
# read in a SupportedMessage response
|
|
header = self.make_header_prefix(SupportedMessage)
|
|
options = self.make_options_body()
|
|
c._socket.recv.return_value = self.make_msg(header, options)
|
|
c.handle_read(None, None)
|
|
|
|
# let it write out a StartupMessage
|
|
c.handle_write(None, None)
|
|
|
|
header = self.make_header_prefix(ServerError, stream_id=1)
|
|
body = self.make_error_body(ServerError.error_code, ServerError.summary)
|
|
c._socket.recv.return_value = self.make_msg(header, body)
|
|
c.handle_read(None, None)
|
|
|
|
# make sure it errored correctly
|
|
self.assertTrue(c.is_defunct)
|
|
self.assertTrue(isinstance(c.last_error, ConnectionException))
|
|
self.assertTrue(c.connected_event.is_set())
|
|
|
|
def test_socket_error_on_write(self, *args):
|
|
c = self.make_connection()
|
|
|
|
# make the OptionsMessage write fail
|
|
c._socket.send.side_effect = socket_error(errno.EIO, "bad stuff!")
|
|
c.handle_write(None, None)
|
|
|
|
# make sure it errored correctly
|
|
self.assertTrue(c.is_defunct)
|
|
self.assertTrue(isinstance(c.last_error, socket_error))
|
|
self.assertTrue(c.connected_event.is_set())
|
|
|
|
def test_blocking_on_write(self, *args):
|
|
c = self.make_connection()
|
|
|
|
# make the OptionsMessage write block
|
|
c._socket.send.side_effect = socket_error(errno.EAGAIN, "socket busy")
|
|
c.handle_write(None, None)
|
|
|
|
self.assertFalse(c.is_defunct)
|
|
|
|
# try again with normal behavior
|
|
c._socket.send.side_effect = lambda x: len(x)
|
|
c.handle_write(None, None)
|
|
self.assertFalse(c.is_defunct)
|
|
self.assertTrue(c._socket.send.call_args is not None)
|
|
|
|
def test_partial_send(self, *args):
|
|
c = self.make_connection()
|
|
|
|
# only write the first four bytes of the OptionsMessage
|
|
c._socket.send.side_effect = None
|
|
c._socket.send.return_value = 4
|
|
c.handle_write(None, None)
|
|
|
|
orig_msg = c._socket.send.call_args[0][0]
|
|
self.assertFalse(c.is_defunct)
|
|
|
|
# try again with normal behavior
|
|
c._socket.send.side_effect = lambda x: len(x)
|
|
c.handle_write(None, None)
|
|
self.assertFalse(c.is_defunct)
|
|
self.assertEqual(c._socket.send.call_args[0][0], orig_msg[4:])
|
|
|
|
def test_socket_error_on_read(self, *args):
|
|
c = self.make_connection()
|
|
|
|
# let it write the OptionsMessage
|
|
c.handle_write(None, None)
|
|
|
|
# read in a SupportedMessage response
|
|
c._socket.recv.side_effect = socket_error(errno.EIO, "busy socket")
|
|
c.handle_read(None, None)
|
|
|
|
# make sure it errored correctly
|
|
self.assertTrue(c.is_defunct)
|
|
self.assertTrue(isinstance(c.last_error, socket_error))
|
|
self.assertTrue(c.connected_event.is_set())
|
|
|
|
def test_partial_header_read(self, *args):
|
|
c = self.make_connection()
|
|
|
|
header = self.make_header_prefix(SupportedMessage)
|
|
options = self.make_options_body()
|
|
message = self.make_msg(header, options)
|
|
|
|
# read in the first byte
|
|
c._socket.recv.return_value = message[0]
|
|
c.handle_read(None, None)
|
|
self.assertEquals(c._buf, message[0])
|
|
|
|
c._socket.recv.return_value = message[1:]
|
|
c.handle_read(None, None)
|
|
self.assertEquals("", c._buf)
|
|
|
|
# let it write out a StartupMessage
|
|
c.handle_write(None, None)
|
|
|
|
header = self.make_header_prefix(ReadyMessage, stream_id=1)
|
|
c._socket.recv.return_value = self.make_msg(header)
|
|
c.handle_read(None, None)
|
|
|
|
self.assertTrue(c.connected_event.is_set())
|
|
self.assertFalse(c.is_defunct)
|
|
|
|
def test_partial_message_read(self, *args):
|
|
c = self.make_connection()
|
|
|
|
header = self.make_header_prefix(SupportedMessage)
|
|
options = self.make_options_body()
|
|
message = self.make_msg(header, options)
|
|
|
|
# read in the first nine bytes
|
|
c._socket.recv.return_value = message[:9]
|
|
c.handle_read(None, None)
|
|
self.assertEquals(c._buf, message[:9])
|
|
|
|
# ... then read in the rest
|
|
c._socket.recv.return_value = message[9:]
|
|
c.handle_read(None, None)
|
|
self.assertEquals("", c._buf)
|
|
|
|
# let it write out a StartupMessage
|
|
c.handle_write(None, None)
|
|
|
|
header = self.make_header_prefix(ReadyMessage, stream_id=1)
|
|
c._socket.recv.return_value = self.make_msg(header)
|
|
c.handle_read(None, None)
|
|
|
|
self.assertTrue(c.connected_event.is_set())
|
|
self.assertFalse(c.is_defunct)
|