Retry socket.timeout as well as SSL timeouts (ssl.SSLError is a subclass of socket.error).

tests/rietveld_test.py

 #!/usr/bin/env python
 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Unit tests for rietveld.py."""
 
 import logging
 import os
+import socket
 import ssl
 import sys
 import time
 import traceback
 import unittest
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 from testing_support.patches_data import GIT, RAW
 from testing_support import auto_stub
@@ skipping 407 matching lines @@
 
 class ProbeException(Exception):
   """Deep-probe a value."""
   value = None
 
   def __init__(self, value):
     super(ProbeException, self).__init__()
     self.value = value
 
 
-def MockSend(request_path, payload=None,
-             content_type="application/octet-stream",
-             timeout=None,
-             extra_headers=None,
-             **kwargs):
+def MockSend(*args, **kwargs):
   """Mock upload.py's Send() to probe the timeout value"""
-  raise ProbeException(timeout)
+  raise ProbeException(kwargs['timeout'])
 
-def MockSendTimeout(request_path, payload=None,
-                    content_type="application/octet-stream",
-                    timeout=None,
-                    extra_headers=None,
-                    **kwargs):
+
+def MockSendTimeout(*args, **kwargs):
   """Mock upload.py's Send() to raise SSLError"""
   raise ssl.SSLError('The read operation timed out')
 
 
+def MockSocketConnectTimeout(*args, **kwargs):
+  """Mock upload.py's Send() to raise socket.timeout"""
+  raise socket.timeout('timed out')
+
+
 class DefaultTimeoutTest(auto_stub.TestCase):
   TESTED_CLASS = rietveld.Rietveld
 
   def setUp(self):
     super(DefaultTimeoutTest, self).setUp()
     self.rietveld = self.TESTED_CLASS('url', None, 'email')
     self.mock(self.rietveld.rpc_server, 'Send', MockSend)
     self.sleep_time = 0
 
   def test_timeout_get(self):
     with self.assertRaises(ProbeException) as cm:
       self.rietveld.get('/api/1234')
 
     self.assertIsNotNone(cm.exception.value, 'Rietveld timeout was not set: %s'
                          % traceback.format_exc())
 
   def test_timeout_post(self):
     with self.assertRaises(ProbeException) as cm:
       self.rietveld.post('/api/1234', [('key', 'data')])
 
     self.assertIsNotNone(cm.exception.value, 'Rietveld timeout was not set: %s'
                          % traceback.format_exc())
 
   def MockSleep(self, t):
     self.sleep_time = t
 
   def test_ssl_timeout_post(self):
     self.mock(self.rietveld.rpc_server, 'Send', MockSendTimeout)
     self.mock(time, 'sleep', self.MockSleep)
-    self.sleep_time = 0
     with self.assertRaises(ssl.SSLError):
       self.rietveld.post('/api/1234', [('key', 'data')])
     self.assertNotEqual(self.sleep_time, 0)
 
+  def test_socket_connect_timeout_post(self):
+    self.mock(self.rietveld.rpc_server, 'Send', MockSocketConnectTimeout)
+    self.mock(time, 'sleep', self.MockSleep)
+    with self.assertRaises(socket.timeout):
+      self.rietveld.post('/api/1234', [('key', 'data')])
+    self.assertNotEqual(self.sleep_time, 0)
+
 if __name__ == '__main__':
   logging.basicConfig(level=[
     logging.ERROR, logging.INFO, logging.DEBUG][min(2, sys.argv.count('-v'))])
   unittest.main()