| OLD | NEW |
| (Empty) |
| 1 import threading | |
| 2 import socket | |
| 3 import select | |
| 4 | |
| 5 | |
| 6 def consume_socket_content(sock, timeout=0.5): | |
| 7 chunks = 65536 | |
| 8 content = b'' | |
| 9 more_to_read = select.select([sock], [], [], timeout)[0] | |
| 10 | |
| 11 while more_to_read: | |
| 12 new_content = sock.recv(chunks) | |
| 13 | |
| 14 if not new_content: | |
| 15 break | |
| 16 | |
| 17 content += new_content | |
| 18 # stop reading if no new data is received for a while | |
| 19 more_to_read = select.select([sock], [], [], timeout)[0] | |
| 20 | |
| 21 return content | |
| 22 | |
| 23 | |
| 24 class Server(threading.Thread): | |
| 25 """Dummy server using for unit testing""" | |
| 26 WAIT_EVENT_TIMEOUT = 5 | |
| 27 | |
| 28 def __init__(self, handler, host='localhost', port=0, requests_to_handle=1,
wait_to_close_event=None): | |
| 29 super(Server, self).__init__() | |
| 30 | |
| 31 self.handler = handler | |
| 32 self.handler_results = [] | |
| 33 | |
| 34 self.host = host | |
| 35 self.port = port | |
| 36 self.requests_to_handle = requests_to_handle | |
| 37 | |
| 38 self.wait_to_close_event = wait_to_close_event | |
| 39 self.ready_event = threading.Event() | |
| 40 self.stop_event = threading.Event() | |
| 41 | |
| 42 @classmethod | |
| 43 def text_response_server(cls, text, request_timeout=0.5, **kwargs): | |
| 44 def text_response_handler(sock): | |
| 45 request_content = consume_socket_content(sock, timeout=request_timeo
ut) | |
| 46 sock.send(text.encode('utf-8')) | |
| 47 | |
| 48 return request_content | |
| 49 | |
| 50 | |
| 51 return Server(text_response_handler, **kwargs) | |
| 52 | |
| 53 @classmethod | |
| 54 def basic_response_server(cls, **kwargs): | |
| 55 return cls.text_response_server( | |
| 56 "HTTP/1.1 200 OK\r\n" + | |
| 57 "Content-Length: 0\r\n\r\n", | |
| 58 **kwargs | |
| 59 ) | |
| 60 | |
| 61 def run(self): | |
| 62 try: | |
| 63 sock = self._create_socket_and_bind() | |
| 64 # in case self.port = 0 | |
| 65 self.port = sock.getsockname()[1] | |
| 66 self.ready_event.set() | |
| 67 self._handle_requests(sock) | |
| 68 | |
| 69 if self.wait_to_close_event: | |
| 70 self.wait_to_close_event.wait(self.WAIT_EVENT_TIMEOUT) | |
| 71 finally: | |
| 72 self.ready_event.set() # just in case of exception | |
| 73 sock.close() | |
| 74 self.stop_event.set() | |
| 75 | |
| 76 def _create_socket_and_bind(self): | |
| 77 sock = socket.socket() | |
| 78 sock.bind((self.host, self.port)) | |
| 79 sock.listen(0) | |
| 80 return sock | |
| 81 | |
| 82 def _handle_requests(self, server_sock): | |
| 83 for _ in range(self.requests_to_handle): | |
| 84 sock = server_sock.accept()[0] | |
| 85 handler_result = self.handler(sock) | |
| 86 | |
| 87 self.handler_results.append(handler_result) | |
| 88 | |
| 89 def __enter__(self): | |
| 90 self.start() | |
| 91 self.ready_event.wait(self.WAIT_EVENT_TIMEOUT) | |
| 92 return self.host, self.port | |
| 93 | |
| 94 def __exit__(self, exc_type, exc_value, traceback): | |
| 95 if exc_type is None: | |
| 96 self.stop_event.wait(self.WAIT_EVENT_TIMEOUT) | |
| 97 else: | |
| 98 if self.wait_to_close_event: | |
| 99 # avoid server from waiting for event timeouts | |
| 100 # if an exception is found in the main thread | |
| 101 self.wait_to_close_event.set() | |
| 102 return False # allow exceptions to propagate | |
| OLD | NEW |