add socket-getting utils

Author: mambocab

tests/unit/io/test_asyncorereactor.py

@@ -61,28 +61,6 @@ def setUp(self):
         if is_monkey_patched():
             raise unittest.SkipTest("Can't test asyncore with monkey patching")
 
-    def test_successful_connection(self):
-        c = self.make_connection()
-
-        # let it write the OptionsMessage
-        c.handle_write()
-
-        # 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()
-
-        # let it write out a StartupMessage
-        c.handle_write()
-
-        header = self.make_header_prefix(ReadyMessage, stream_id=1)
-        c.socket.recv.return_value = self.make_msg(header)
-        c.handle_read()
-
-        self.assertTrue(c.connected_event.is_set())
-        return c
-
     def test_egain_on_buffer_size(self):
         # get a connection that's already fully started
         c = self.test_successful_connection()

tests/unit/io/test_libevreactor.py

@@ -43,6 +43,7 @@ class LibevConnectionTest(unittest.TestCase, ReactorTestMixin):
 
     connection_class = LibevConnection
     socket_attr_name = '_socket'
+    null_handle_function_args = None, 0
 
     def setUp(self):
         if is_monkey_patched():
@@ -64,28 +65,6 @@ def setUp(self):
         for p in patchers:
             p.start()
 
-    def test_successful_connection(self):
-        c = self.make_connection()
-
-        # let it write the OptionsMessage
-        c.handle_write(None, 0)
-
-        # 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, 0)
-
-        # let it write out a StartupMessage
-        c.handle_write(None, 0)
-
-        header = self.make_header_prefix(ReadyMessage, stream_id=1)
-        c._socket.recv.return_value = self.make_msg(header)
-        c.handle_read(None, 0)
-
-        self.assertTrue(c.connected_event.is_set())
-        return c
-
     def test_egain_on_buffer_size(self):
         # get a connection that's already fully started
         c = self.test_successful_connection()

tests/unit/io/utils.py

@@ -14,7 +14,8 @@
 
 from cassandra.connection import HEADER_DIRECTION_TO_CLIENT
 from cassandra.marshal import uint8_pack, uint32_pack
-from cassandra.protocol import write_stringmultimap, write_int, write_string
+from cassandra.protocol import (write_stringmultimap, write_int, write_string,
+                                SupportedMessage, ReadyMessage)
 
 from six import binary_type, BytesIO
 from mock import Mock
@@ -151,6 +152,13 @@ def test_timer_cancellation(self):
 class ReactorTestMixin(object):
 
     connection_class = socket_attr_name = None
+    null_handle_function_args = ()
+
+    def get_socket(self, connection):
+        return getattr(connection, self.socket_attr_name)
+
+    def set_socket(self, connection, obj):
+        return setattr(connection, self.socket_attr_name, obj)
 
     def make_header_prefix(self, message_class, version=2, stream_id=0):
         return binary_type().join(map(uint8_pack, [
@@ -164,7 +172,7 @@ def make_connection(self):
         c = self.connection_class('1.2.3.4', cql_version='3.0.1', connect_timeout=5)
         mocket = Mock()
         mocket.send.side_effect = lambda x: len(x)
-        setattr(c, self.socket_attr_name, mocket)
+        self.set_socket(c, mocket)
         return c
 
     def make_options_body(self):
@@ -183,3 +191,25 @@ def make_error_body(self, code, msg):
 
     def make_msg(self, header, body=binary_type()):
         return header + uint32_pack(len(body)) + body
+
+    def test_successful_connection(self):
+        c = self.make_connection()
+
+        # let it write the OptionsMessage
+        c.handle_write(*self.null_handle_function_args)
+
+        # read in a SupportedMessage response
+        header = self.make_header_prefix(SupportedMessage)
+        options = self.make_options_body()
+        self.get_socket(c).recv.return_value = self.make_msg(header, options)
+        c.handle_read(*self.null_handle_function_args)
+
+        # let it write out a StartupMessage
+        c.handle_write(*self.null_handle_function_args)
+
+        header = self.make_header_prefix(ReadyMessage, stream_id=1)
+        self.get_socket(c).recv.return_value = self.make_msg(header)
+        c.handle_read(*self.null_handle_function_args)
+
+        self.assertTrue(c.connected_event.is_set())
+        return c