|
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") |