Improve multi_callback_test

1 files changed, 26 insertions, 38 deletions

diff --git a/tests/multi_callback_test.py b/tests/multi_callback_test.py
index c82bc38..a86cb43 100644
--- a/tests/multi_callback_test.py
+++ b/tests/multi_callback_test.py
@@ -14,13 +14,13 @@ setup_module, teardown_module = appmanager.setup(('app', 8380))
 class MultiCallbackTest(unittest.TestCase):
     def setUp(self):
         self.easy = util.DefaultCurl()
-        self.easy.setopt(pycurl.URL, 'http://%s:8380/success' % localhost)
+        self.easy.setopt(pycurl.URL, 'http://%s:8380/long_pause' % localhost)
         self.multi = pycurl.CurlMulti()
         self.multi.setopt(pycurl.M_SOCKETFUNCTION, self.socket_callback)
         self.multi.setopt(pycurl.M_TIMERFUNCTION, self.timer_callback)
-        self.timer_result = None
         self.socket_result = None
-        self.socket_action = None
+        self.timer_result = None
+        self.sockets = {}
 
     def tearDown(self):
         self.multi.close()
@@ -29,13 +29,24 @@ class MultiCallbackTest(unittest.TestCase):
     def socket_callback(self, ev_bitmask, sock_fd, multi, data):
         self.socket_result = (sock_fd, ev_bitmask)
         if ev_bitmask & pycurl.POLL_REMOVE:
-            pass
+            self.sockets.pop(sock_fd)
         else:
-            self.socket_action = (sock_fd, ev_bitmask)
+            self.sockets[sock_fd] = ev_bitmask | self.sockets.get(sock_fd, 0)
 
     def timer_callback(self, timeout_ms):
         self.timer_result = timeout_ms
-        self.socket_action = (pycurl.SOCKET_TIMEOUT, 0)
+
+    def partial_transfer(self):
+        perform = True
+        def write_callback(data):
+            nonlocal perform
+            perform = False
+        self.easy.setopt(pycurl.WRITEFUNCTION, write_callback)
+        self.multi.add_handle(self.easy)
+        self.multi.socket_action(pycurl.SOCKET_TIMEOUT, 0)
+        while self.sockets and perform:
+            for socket, action in tuple(self.sockets.items()):
+                self.multi.socket_action(socket, action)
 
     # multi.socket_action must call both SOCKETFUNCTION and TIMERFUNCTION at
     # various points during the transfer (at least at the start and end)
@@ -43,64 +54,41 @@ class MultiCallbackTest(unittest.TestCase):
         self.multi.add_handle(self.easy)
         self.timer_result = None
         self.socket_result = None
-        while self.multi.socket_action(*self.socket_action)[1]:
-            # Without real event loop we just use blocking select call instead
-            self.multi.select(0.1)
-        # both callbacks should be invoked multiple times by socket_action
+        self.multi.socket_action(pycurl.SOCKET_TIMEOUT, 0)
         assert self.socket_result is not None
         assert self.timer_result is not None
 
     # multi.add_handle must call TIMERFUNCTION to schedule a kick-start
     def test_multi_add_handle(self):
-        assert self.timer_result == None
         self.multi.add_handle(self.easy)
         assert self.timer_result is not None
 
     # (mid-transfer) multi.remove_handle must call SOCKETFUNCTION to remove sockets
     def test_multi_remove_handle(self):
         self.multi.add_handle(self.easy)
-        while self.multi.socket_action(*self.socket_action)[1]:
-            if self.socket_result:
-                # libcurl informed us about new sockets
-                break
-            # Without real event loop we just use blocking select call instead
-            self.multi.select(0.1)
+        self.multi.socket_action(pycurl.SOCKET_TIMEOUT, 0)
         self.socket_result = None
-        # libcurl will now inform us that we should remove those sockets
         self.multi.remove_handle(self.easy)
         assert self.socket_result is not None
-    
+
     # (mid-transfer) easy.pause(PAUSE_ALL) must call SOCKETFUNCTION to remove sockets
     # (mid-transfer) easy.pause(PAUSE_CONT) must call TIMERFUNCTION to resume
     def test_easy_pause_unpause(self):
-        self.multi.add_handle(self.easy)
-        while self.multi.socket_action(*self.socket_action)[1]:
-            if self.socket_result:
-                # libcurl informed us about new sockets
-                break
-            # Without real event loop we just use blocking select call instead
-            self.multi.select(0.1)
+        self.partial_transfer()
         self.socket_result = None
-        # libcurl will now inform us that we should remove those sockets
+        # libcurl will now inform us that we should remove some sockets
         self.easy.pause(pycurl.PAUSE_ALL)
         assert self.socket_result is not None
+        self.socket_result = None
         self.timer_result = None
-        # libcurl will now tell us to schedule a kickstart
+        # libcurl will now tell us to add those sockets and schedule a kickstart
         self.easy.pause(pycurl.PAUSE_CONT)
+        assert self.socket_result is not None
         assert self.timer_result is not None
 
     # (mid-transfer) easy.close() must call SOCKETFUNCTION to remove sockets
-    # NOTE: doing easy.close() during transfer is considered wrong,
-    # but libcurl still invokes callbacks to inform your event loop
     def test_easy_close(self):
-        self.multi.add_handle(self.easy)
-        while self.multi.socket_action(*self.socket_action)[1]:
-            if self.socket_result:
-                # libcurl informed us about new sockets
-                break
-            # Without real event loop we just use blocking select call instead
-            self.multi.select(0.1)
+        self.partial_transfer()
         self.socket_result = None
-        # libcurl will now inform us that we should remove those sockets
         self.easy.close()
         assert self.socket_result is not None