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         data = self.rfile.read(int(self.headers.getheader('content-length')))
  69         datahash = hashlib.md5(data).hexdigest()
  70         self.server.store[datahash] = data
  71         self.server._do_delay('response')
  72         self.send_response(200)
  73         self.send_header('Content-type', 'text/plain')
  74         self.end_headers()
  75         self.server._do_delay('response_body')
  76         self.wfile.write(datahash + '+' + str(len(data)))
  77         self.server._do_delay('response_close')
  78
  79     def log_request(self, *args, **kwargs):
  80         if os.environ.get('ARVADOS_DEBUG', None):
  81             super(Handler, self).log_request(*args, **kwargs)
  82
  83     def finish(self, *args, **kwargs):
  84         """Ignore exceptions, notably "Broken pipe" when client times out."""
  85         try:
  86             return super(Handler, self).finish(*args, **kwargs)
  87         except:
  88             pass
  89
  90     def handle_one_request(self, *args, **kwargs):
  91         """Ignore exceptions, notably "Broken pipe" when client times out."""
  92         try:
  93             return super(Handler, self).handle_one_request(*args, **kwargs)
  94         except:
  95             pass