Fix issue 9794: adds context manager protocol to socket.socket so that socket.create_connection() can be used with the 'with' statement.
diff --git a/Lib/test/test_socket.py b/Lib/test/test_socket.py
index 19c494b..81f9cdf 100644
--- a/Lib/test/test_socket.py
+++ b/Lib/test/test_socket.py
@@ -1595,6 +1595,49 @@
self.cli.close()
+@unittest.skipUnless(thread, 'Threading required for this test.')
+class ContextManagersTest(ThreadedTCPSocketTest):
+
+ def _testSocketClass(self):
+ # base test
+ with socket.socket() as sock:
+ self.assertFalse(sock._closed)
+ self.assertTrue(sock._closed)
+ # close inside with block
+ with socket.socket() as sock:
+ sock.close()
+ self.assertTrue(sock._closed)
+ # exception inside with block
+ with socket.socket() as sock:
+ self.assertRaises(socket.error, sock.sendall, b'foo')
+ self.assertTrue(sock._closed)
+
+ def testCreateConnectionBase(self):
+ conn, addr = self.serv.accept()
+ data = conn.recv(1024)
+ conn.sendall(data)
+
+ def _testCreateConnectionBase(self):
+ address = self.serv.getsockname()
+ with socket.create_connection(address) as sock:
+ self.assertFalse(sock._closed)
+ sock.sendall(b'foo')
+ self.assertEqual(sock.recv(1024), b'foo')
+ self.assertTrue(sock._closed)
+
+ def testCreateConnectionClose(self):
+ conn, addr = self.serv.accept()
+ data = conn.recv(1024)
+ conn.sendall(data)
+
+ def _testCreateConnectionClose(self):
+ address = self.serv.getsockname()
+ with socket.create_connection(address) as sock:
+ sock.close()
+ self.assertTrue(sock._closed)
+ self.assertRaises(socket.error, sock.sendall, b'foo')
+
+
def test_main():
tests = [GeneralModuleTests, BasicTCPTest, TCPCloserTest, TCPTimeoutTest,
TestExceptions, BufferIOTest, BasicTCPTest2, BasicUDPTest, UDPTimeoutTest ]
@@ -1609,6 +1652,7 @@
NetworkConnectionNoServer,
NetworkConnectionAttributesTest,
NetworkConnectionBehaviourTest,
+ ContextManagersTest,
])
if hasattr(socket, "socketpair"):
tests.append(BasicSocketPairTest)