1 files changed, 113 insertions, 10 deletions
diff --git a/Lib/test/test_nntplib.py b/Lib/test/test_nntplib.py
index 1d52713..ae3618f 100644
--- a/Lib/test/test_nntplib.py
+++ b/Lib/test/test_nntplib.py
@@ -8,6 +8,7 @@ import contextlib
 from test import support
 from nntplib import NNTP, GroupInfo
 import nntplib
+from unittest.mock import patch
 try:
     import ssl
 except ImportError:
@@ -266,7 +267,7 @@ class NetworkedNNTPTestsMixin:
                 return False
             try:
                 server.help()
-            except (socket.error, EOFError):
+            except (OSError, EOFError):
                 return False
             return True
 
@@ -370,6 +371,14 @@ class _NNTPServerIO(io.RawIOBase):
         return n
 
 
+def make_mock_file(handler):
+    sio = _NNTPServerIO(handler)
+    # Using BufferedRWPair instead of BufferedRandom ensures the file
+    # isn't seekable.
+    file = io.BufferedRWPair(sio, sio)
+    return (sio, file)
+
+
 class MockedNNTPTestsMixin:
     # Override in derived classes
     handler_class = None
@@ -384,10 +393,7 @@ class MockedNNTPTestsMixin:
 
     def make_server(self, *args, **kwargs):
         self.handler = self.handler_class()
-        self.sio = _NNTPServerIO(self.handler)
-        # Using BufferedRWPair instead of BufferedRandom ensures the file
-        # isn't seekable.
-        file = io.BufferedRWPair(self.sio, self.sio)
+        self.sio, file = make_mock_file(self.handler)
         self.server = nntplib._NNTPBase(file, 'test.server', *args, **kwargs)
         return self.server
 
@@ -1412,11 +1418,108 @@ class MiscTests(unittest.TestCase):
     def test_ssl_support(self):
         self.assertTrue(hasattr(nntplib, 'NNTP_SSL'))
 
-def test_main():
-    tests = [MiscTests, NNTPv1Tests, NNTPv2Tests, CapsAfterLoginNNTPv2Tests,
-            SendReaderNNTPv2Tests, NetworkedNNTPTests, NetworkedNNTP_SSLTests]
-    support.run_unittest(*tests)
+
+class PublicAPITests(unittest.TestCase):
+    """Ensures that the correct values are exposed in the public API."""
+
+    def test_module_all_attribute(self):
+        self.assertTrue(hasattr(nntplib, '__all__'))
+        target_api = ['NNTP', 'NNTPError', 'NNTPReplyError',
+                      'NNTPTemporaryError', 'NNTPPermanentError',
+                      'NNTPProtocolError', 'NNTPDataError', 'decode_header']
+        if ssl is not None:
+            target_api.append('NNTP_SSL')
+        self.assertEqual(set(nntplib.__all__), set(target_api))
+
+class MockSocketTests(unittest.TestCase):
+    """Tests involving a mock socket object
+
+    Used where the _NNTPServerIO file object is not enough."""
+
+    nntp_class = nntplib.NNTP
+
+    def check_constructor_error_conditions(
+            self, handler_class,
+            expected_error_type, expected_error_msg,
+            login=None, password=None):
+
+        class mock_socket_module:
+            def create_connection(address, timeout):
+                return MockSocket()
+
+        class MockSocket:
+            def close(self):
+                nonlocal socket_closed
+                socket_closed = True
+
+            def makefile(socket, mode):
+                handler = handler_class()
+                _, file = make_mock_file(handler)
+                files.append(file)
+                return file
+
+        socket_closed = False
+        files = []
+        with patch('nntplib.socket', mock_socket_module), \
+             self.assertRaisesRegex(expected_error_type, expected_error_msg):
+            self.nntp_class('dummy', user=login, password=password)
+        self.assertTrue(socket_closed)
+        for f in files:
+            self.assertTrue(f.closed)
+
+    def test_bad_welcome(self):
+        #Test a bad welcome message
+        class Handler(NNTPv1Handler):
+            welcome = 'Bad Welcome'
+        self.check_constructor_error_conditions(
+            Handler, nntplib.NNTPProtocolError, Handler.welcome)
+
+    def test_service_temporarily_unavailable(self):
+        #Test service temporarily unavailable
+        class Handler(NNTPv1Handler):
+            welcome = '400 Service temporarily unavilable'
+        self.check_constructor_error_conditions(
+            Handler, nntplib.NNTPTemporaryError, Handler.welcome)
+
+    def test_service_permanently_unavailable(self):
+        #Test service permanently unavailable
+        class Handler(NNTPv1Handler):
+            welcome = '502 Service permanently unavilable'
+        self.check_constructor_error_conditions(
+            Handler, nntplib.NNTPPermanentError, Handler.welcome)
+
+    def test_bad_capabilities(self):
+        #Test a bad capabilities response
+        class Handler(NNTPv1Handler):
+            def handle_CAPABILITIES(self):
+                self.push_lit(capabilities_response)
+        capabilities_response = '201 bad capability'
+        self.check_constructor_error_conditions(
+            Handler, nntplib.NNTPReplyError, capabilities_response)
+
+    def test_login_aborted(self):
+        #Test a bad authinfo response
+        login = 't@e.com'
+        password = 'python'
+        class Handler(NNTPv1Handler):
+            def handle_AUTHINFO(self, *args):
+                self.push_lit(authinfo_response)
+        authinfo_response = '503 Mechanism not recognized'
+        self.check_constructor_error_conditions(
+            Handler, nntplib.NNTPPermanentError, authinfo_response,
+            login, password)
+
+class bypass_context:
+    """Bypass encryption and actual SSL module"""
+    def wrap_socket(sock, **args):
+        return sock
+
+@unittest.skipUnless(ssl, 'requires SSL support')
+class MockSslTests(MockSocketTests):
+    @staticmethod
+    def nntp_class(*pos, **kw):
+        return nntplib.NNTP_SSL(*pos, ssl_context=bypass_context, **kw)
 
 
 if __name__ == "__main__":
-    test_main()
+    unittest.main()