sdk/python/tests/keepstub.py

   1 import BaseHTTPServer
   2 import hashlib
   3 import os
   4 import re
   5 import SocketServer
   6 import time
   7
   8 class Server(SocketServer.ThreadingMixIn, BaseHTTPServer.HTTPServer, object):
   9
  10     allow_reuse_address = 1
  11
  12     def __init__(self, *args, **kwargs):
  13         self.store = {}
  14         self.delays = {
  15             # before reading request headers
  16             'request': 0,
  17             # before reading request body
  18             'request_body': 0,
  19             # before setting response status and headers
  20             'response': 0,
  21             # before sending response body
  22             'response_body': 0,
  23             # before returning from handler (thus setting response EOF)
  24             'response_close': 0,
  25         }
  26         super(Server, self).__init__(*args, **kwargs)
  27
  28     def setdelays(self, **kwargs):
  29         """In future requests, induce delays at the given checkpoints."""
  30         for (k, v) in kwargs.iteritems():
  31             self.delays.get(k) # NameError if unknown key
  32             self.delays[k] = v
  33
  34     def _sleep_at_least(self, seconds):
  35         """Sleep for given time, even if signals are received."""
  36         wake = time.time() + seconds
  37         todo = seconds
  38         while todo > 0:
  39             time.sleep(todo)
  40             todo = wake - time.time()
  41
  42     def _do_delay(self, k):
  43         self._sleep_at_least(self.delays[k])
  44
  45
  46 class Handler(BaseHTTPServer.BaseHTTPRequestHandler, object):
  47     def handle(self, *args, **kwargs):
  48         self.server._do_delay('request')
  49         return super(Handler, self).handle(*args, **kwargs)
  50
  51     def do_GET(self):
  52         self.server._do_delay('response')
  53         r = re.search(r'[0-9a-f]{32}', self.path)
  54         if not r:
  55             return self.send_response(422)
  56         datahash = r.group(0)
  57         if datahash not in self.server.store:
  58             return self.send_response(404)
  59         self.send_response(200)
  60         self.send_header('Content-type', 'application/octet-stream')
  61         self.end_headers()
  62         self.server._do_delay('response_body')
  63         self.wfile.write(self.server.store[datahash])
  64         self.server._do_delay('response_close')
  65
  66     def do_PUT(self):
  67         self.server._do_delay('request_body')
  68
  69         # The comments at https://bugs.python.org/issue1491 implies that Python
  70         # 2.7 BaseHTTPRequestHandler was patched to support 100 Continue, but
  71         # reading the actual code that ships in Debian it clearly is not, so we
  72         # need to send the response on the socket directly.
  73
  74         self.wfile.write("%s %d %s\r\n\r\n" %
  75                          (self.protocol_version, 100, "Continue"))
  76
  77         data = self.rfile.read(int(self.headers.getheader('content-length')))
  78         datahash = hashlib.md5(data).hexdigest()
  79         self.server.store[datahash] = data
  80         self.server._do_delay('response')
  81         self.send_response(200)
  82         self.send_header('Content-type', 'text/plain')
  83         self.end_headers()
  84         self.server._do_delay('response_body')
  85         self.wfile.write(datahash + '+' + str(len(data)))
  86         self.server._do_delay('response_close')
  87
  88     def log_request(self, *args, **kwargs):
  89         if os.environ.get('ARVADOS_DEBUG', None):
  90             super(Handler, self).log_request(*args, **kwargs)
  91
  92     def finish(self, *args, **kwargs):
  93         """Ignore exceptions, notably "Broken pipe" when client times out."""
  94         try:
  95             return super(Handler, self).finish(*args, **kwargs)
  96         except:
  97             pass
  98
  99     def handle_one_request(self, *args, **kwargs):
 100         """Ignore exceptions, notably "Broken pipe" when client times out."""
 101         try:
 102             return super(Handler, self).handle_one_request(*args, **kwargs)
 103         except:
 104             pass