Browse Source

[1828] Make code safer by using with (for sock)

Mukund Sivaraman 13 years ago
parent
commit
ff9dd7da09
1 changed files with 57 additions and 58 deletions
  1. 57 58
      src/lib/python/isc/util/cio/tests/socketsession_test.py

+ 57 - 58
src/lib/python/isc/util/cio/tests/socketsession_test.py

@@ -126,64 +126,63 @@ class TestForwarder(unittest.TestCase):
 
 
     def check_push_and_pop(self, family, type, protocol, local, remote,
     def check_push_and_pop(self, family, type, protocol, local, remote,
                            data, new_connection):
                            data, new_connection):
-        sock = self.create_socket(family, type, protocol, local, True)
-        fwd_fd = sock.fileno()
-        if protocol == IPPROTO_TCP:
-            client_addr = ('::1', 0, 0, 0) if family == AF_INET6 \
-                else ('127.0.0.1', 0)
-            client_sock = self.create_socket(family, type, protocol,
-                                             client_addr, False)
-            client_sock.setblocking(False)
-            try:
-                client_sock.connect(local)
-            except socket.error:
-                pass
-            server_sock, _ = sock.accept()
-            fwd_fd = server_sock.fileno()
-
-        # If a new connection is required, start the "server", have the
-        # internal forwarder connect to it, and then internally accept it.
-        if new_connection:
-            self.start_listen()
-            self.forwarder.connect_to_receiver()
-            self.accept_sock = self.accept_forwarder()
-
-        # Then push one socket session via the forwarder.
-        self.forwarder.push(fwd_fd, family, type, protocol, local, remote,
-                            data)
-
-        # Pop the socket session we just pushed from a local receiver, and
-        # check the content.
-        receiver = SocketSessionReceiver(self.accept_sock)
-        signal.alarm(1)
-        sock_session = receiver.pop()
-        signal.alarm(0)
-        passed_sock = sock_session[0]
-        self.assertNotEqual(fwd_fd, passed_sock.fileno())
-        self.assertEqual(family, passed_sock.family)
-        self.assertEqual(type, passed_sock.type)
-        self.assertEqual(protocol, passed_sock.proto)
-        self.assertEqual(local, sock_session[1])
-        self.assertEqual(remote, sock_session[2])
-        self.assertEqual(data, sock_session[3])
-
-        # Check if the passed FD is usable by sending some data from it.
-        passed_sock.setblocking(True)
-        if protocol == IPPROTO_UDP:
-            self.assertEqual(len(TEST_DATA), passed_sock.sendto(TEST_DATA,
-                                                                local))
-            sock.settimeout(10)
-            self.assertEqual(TEST_DATA, sock.recvfrom(len(TEST_DATA))[0])
-        else:
-            self.assertEqual(len(TEST_DATA), passed_sock.send(TEST_DATA))
-            client_sock.setblocking(True)
-            client_sock.settimeout(10)
-            self.assertEqual(TEST_DATA, client_sock.recv(len(TEST_DATA)))
-            server_sock.close()
-            client_sock.close()
-
-        passed_sock.close()
-        sock.close()
+        with self.create_socket(family, type, protocol, local, True) as sock:
+            fwd_fd = sock.fileno()
+            if protocol == IPPROTO_TCP:
+                client_addr = ('::1', 0, 0, 0) if family == AF_INET6 \
+                    else ('127.0.0.1', 0)
+                client_sock = self.create_socket(family, type, protocol,
+                                                 client_addr, False)
+                client_sock.setblocking(False)
+                try:
+                    client_sock.connect(local)
+                except socket.error:
+                    pass
+                server_sock, _ = sock.accept()
+                fwd_fd = server_sock.fileno()
+
+            # If a new connection is required, start the "server", have the
+            # internal forwarder connect to it, and then internally accept it.
+            if new_connection:
+                self.start_listen()
+                self.forwarder.connect_to_receiver()
+                self.accept_sock = self.accept_forwarder()
+
+            # Then push one socket session via the forwarder.
+            self.forwarder.push(fwd_fd, family, type, protocol, local, remote,
+                                data)
+
+            # Pop the socket session we just pushed from a local receiver, and
+            # check the content.
+            receiver = SocketSessionReceiver(self.accept_sock)
+            signal.alarm(1)
+            sock_session = receiver.pop()
+            signal.alarm(0)
+            passed_sock = sock_session[0]
+            self.assertNotEqual(fwd_fd, passed_sock.fileno())
+            self.assertEqual(family, passed_sock.family)
+            self.assertEqual(type, passed_sock.type)
+            self.assertEqual(protocol, passed_sock.proto)
+            self.assertEqual(local, sock_session[1])
+            self.assertEqual(remote, sock_session[2])
+            self.assertEqual(data, sock_session[3])
+
+            # Check if the passed FD is usable by sending some data from it.
+            passed_sock.setblocking(True)
+            if protocol == IPPROTO_UDP:
+                self.assertEqual(len(TEST_DATA), passed_sock.sendto(TEST_DATA,
+                                                                    local))
+                sock.settimeout(10)
+                self.assertEqual(TEST_DATA, sock.recvfrom(len(TEST_DATA))[0])
+            else:
+                self.assertEqual(len(TEST_DATA), passed_sock.send(TEST_DATA))
+                client_sock.setblocking(True)
+                client_sock.settimeout(10)
+                self.assertEqual(TEST_DATA, client_sock.recv(len(TEST_DATA)))
+                server_sock.close()
+                client_sock.close()
+
+            passed_sock.close()
 
 
     def test_push_and_pop(self):
     def test_push_and_pop(self):
         # This is a straightforward port of C++ pushAndPop test.
         # This is a straightforward port of C++ pushAndPop test.