script/lib/tweetstream/servercontext.py
changeset 12 4daf47fcf792
equal deleted inserted replaced
11:54d7f1486ac4 12:4daf47fcf792
       
     1 import threading
       
     2 import contextlib
       
     3 import time
       
     4 import os
       
     5 import socket
       
     6 import random
       
     7 from functools import partial
       
     8 from inspect import isclass
       
     9 from BaseHTTPServer import HTTPServer, BaseHTTPRequestHandler
       
    10 from SimpleHTTPServer import SimpleHTTPRequestHandler
       
    11 from SocketServer import BaseRequestHandler
       
    12 
       
    13 
       
    14 class ServerError(Exception):
       
    15     pass
       
    16 
       
    17 
       
    18 class ServerContext(object):
       
    19     """Context object with information about a running test server."""
       
    20 
       
    21     def __init__(self, address, port):
       
    22         self.address = address or "localhost"
       
    23         self.port = port
       
    24 
       
    25     @property
       
    26     def baseurl(self):
       
    27         return "http://%s:%s" % (self.address, self.port)
       
    28 
       
    29     def __str__(self):
       
    30         return "<ServerContext %s >" % self.baseurl
       
    31 
       
    32     __repr__ = __str__
       
    33 
       
    34 
       
    35 class _SilentSimpleHTTPRequestHandler(SimpleHTTPRequestHandler):
       
    36 
       
    37     def __init__(self, *args, **kwargs):
       
    38         self.logging = kwargs.get("logging", False)
       
    39         SimpleHTTPRequestHandler.__init__(self, *args, **kwargs)
       
    40 
       
    41     def log_message(self, *args, **kwargs):
       
    42         if self.logging:
       
    43             SimpleHTTPRequestHandler.log_message(self, *args, **kwargs)
       
    44 
       
    45 
       
    46 class _TestHandler(BaseHTTPRequestHandler):
       
    47     """RequestHandler class that handles requests that use a custom handler
       
    48     callable."""
       
    49 
       
    50     def __init__(self, handler, methods, *args, **kwargs):
       
    51         self._handler = handler
       
    52         self._methods = methods
       
    53         self._response_sent = False
       
    54         self._headers_sent = False
       
    55         self.logging = kwargs.get("logging", False)
       
    56         BaseHTTPRequestHandler.__init__(self, *args, **kwargs)
       
    57 
       
    58     def log_message(self, *args, **kwargs):
       
    59         if self.logging:
       
    60             BaseHTTPRequestHandler.log_message(self, *args, **kwargs)
       
    61 
       
    62     def send_response(self, *args, **kwargs):
       
    63         self._response_sent = True
       
    64         BaseHTTPRequestHandler.send_response(self, *args, **kwargs)
       
    65 
       
    66     def end_headers(self, *args, **kwargs):
       
    67         self._headers_sent = True
       
    68         BaseHTTPRequestHandler.end_headers(self, *args, **kwargs)
       
    69 
       
    70     def _do_whatever(self):
       
    71         """Called in place of do_METHOD"""
       
    72         data = self._handler(self)
       
    73 
       
    74         if hasattr(data, "next"):
       
    75             # assume it's something supporting generator protocol
       
    76             self._handle_with_iterator(data)
       
    77         else:
       
    78             # Nothing more to do then.
       
    79             pass
       
    80 
       
    81 
       
    82     def __getattr__(self, name):
       
    83         if name.startswith("do_") and name[3:].lower() in self._methods:
       
    84             return self._do_whatever
       
    85         else:
       
    86             # fixme instance or class?
       
    87             raise AttributeError(name)
       
    88 
       
    89     def _handle_with_iterator(self, iterator):
       
    90         self.connection.settimeout(0.1)
       
    91         for data in iterator:
       
    92             if not self.server.server_thread.running:
       
    93                 return
       
    94 
       
    95             if not self._response_sent:
       
    96                 self.send_response(200)
       
    97             if not self._headers_sent:
       
    98                 self.end_headers()
       
    99 
       
   100             self.wfile.write(data)
       
   101             # flush immediatly. We may want to do trickling writes
       
   102             # or something else tha trequires bypassing normal caching
       
   103             self.wfile.flush()
       
   104 
       
   105 class _TestServerThread(threading.Thread):
       
   106     """Thread class for a running test server"""
       
   107 
       
   108     def __init__(self, handler, methods, cwd, port, address):
       
   109         threading.Thread.__init__(self)
       
   110         self.startup_finished = threading.Event()
       
   111         self._methods = methods
       
   112         self._cwd = cwd
       
   113         self._orig_cwd = None
       
   114         self._handler = self._wrap_handler(handler, methods)
       
   115         self._setup()
       
   116         self.running = True
       
   117         self.serverloc = (address, port)
       
   118         self.error = None
       
   119 
       
   120     def _wrap_handler(self, handler, methods):
       
   121         if isclass(handler) and issubclass(handler, BaseRequestHandler):
       
   122             return handler # It's OK. user passed in a proper handler
       
   123         elif callable(handler):
       
   124             return partial(_TestHandler, handler, methods)
       
   125             # it's a callable, so wrap in a req handler
       
   126         else:
       
   127             raise ServerError("handler must be callable or RequestHandler")
       
   128 
       
   129     def _setup(self):
       
   130         if self._cwd != "./":
       
   131             self._orig_cwd = os.getcwd()
       
   132             os.chdir(self._cwd)
       
   133 
       
   134     def _init_server(self):
       
   135         """Hooks up the server socket"""
       
   136         try:
       
   137             if self.serverloc[1] == "random":
       
   138                 retries = 10 # try getting an available port max this many times
       
   139                 while True:
       
   140                     try:
       
   141                         self.serverloc = (self.serverloc[0],
       
   142                                           random.randint(1025, 49151))
       
   143                         self._server = HTTPServer(self.serverloc, self._handler)
       
   144                     except socket.error:
       
   145                         retries -= 1
       
   146                         if not retries: # not able to get a port.
       
   147                             raise
       
   148                     else:
       
   149                         break
       
   150             else: # use specific port. this might throw, that's expected
       
   151                 self._server = HTTPServer(self.serverloc, self._handler)
       
   152         except socket.error, e:
       
   153             self.running = False
       
   154             self.error = e
       
   155             # set this here, since we'll never enter the serve loop where
       
   156             # it is usually set:
       
   157             self.startup_finished.set()
       
   158             return
       
   159 
       
   160         self._server.allow_reuse_address = True # lots of tests, same port
       
   161         self._server.timeout = 0.1
       
   162         self._server.server_thread = self
       
   163 
       
   164 
       
   165     def run(self):
       
   166         self._init_server()
       
   167 
       
   168         while self.running:
       
   169             self._server.handle_request() # blocks for self.timeout secs
       
   170             # First time this falls through, signal the parent thread that
       
   171             # the server is ready for incomming connections
       
   172             if not self.startup_finished.is_set():
       
   173                 self.startup_finished.set()
       
   174 
       
   175         self._cleanup()
       
   176 
       
   177     def stop(self):
       
   178         """Stop the server and attempt to make the thread terminate.
       
   179         This happens async but the calling code can check periodically
       
   180         the isRunning flag on the thread object.
       
   181         """
       
   182         # actual stopping happens in the run method
       
   183         self.running = False
       
   184 
       
   185     def _cleanup(self):
       
   186         """Do some rudimentary cleanup."""
       
   187         if self._orig_cwd:
       
   188             os.chdir(self._orig_cwd)
       
   189 
       
   190 
       
   191 @contextlib.contextmanager
       
   192 def test_server(handler=_SilentSimpleHTTPRequestHandler, port=8514,
       
   193                 address="", methods=("get", "head"), cwd="./"):
       
   194     """Context that makes available a web server in a separate thread"""
       
   195     thread = _TestServerThread(handler=handler, methods=methods, cwd=cwd,
       
   196                                port=port, address=address)
       
   197     thread.start()
       
   198 
       
   199     # fixme: should this be daemonized? If it isn't it will block the entire
       
   200     # app, but that should never happen anyway..
       
   201     thread.startup_finished.wait()
       
   202 
       
   203     if thread.error: # startup failed! Bail, throw whatever the server did
       
   204         raise thread.error
       
   205 
       
   206     exc = None
       
   207     try:
       
   208         yield ServerContext(*thread.serverloc)
       
   209     except Exception, exc:
       
   210         pass
       
   211     thread.stop()
       
   212     thread.join(5) # giving it a lot of leeway. should never happen
       
   213 
       
   214     if exc:
       
   215         raise exc
       
   216 
       
   217     # fixme: this takes second priorty after the internal exception but would
       
   218     # still be nice to signal back to calling code.
       
   219 
       
   220     if thread.isAlive():
       
   221         raise Warning("Test server could not be stopped")