From 6551c290fdbf0b29e4534abef1b04f7485a88aad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Braulio=20Valdivielso=20Mart=C3=ADnez?=
 <brlvldvlsmrtnz@gmail.com>
Date: Wed, 2 Dec 2015 01:52:32 +0100
Subject: [PATCH] Allow main thread to block testserver before closing

---
 testserver/server.py | 32 ++++++++++++++++++++++++--------
 1 file changed, 24 insertions(+), 8 deletions(-)

diff --git a/testserver/server.py b/testserver/server.py
index e329a9fc..5c5d8069 100644
--- a/testserver/server.py
+++ b/testserver/server.py
@@ -1,39 +1,47 @@
 #!/usr/bin/python
 
-import threading, socket
+import threading 
+import socket
+
+def consume_socket(sock, chunks=65536):
+    while not sock.recv(chunks).endswith(b'\r\n\r\n'):
+            pass
+
 
 class Server(threading.Thread):
     """ Dummy server using for unit testing """
 
-    def __init__(self, handler, host='localhost', port=0):
+    def __init__(self, handler, host='localhost', port=0, requests_to_handle=1, wait_to_close_event=None):
         threading.Thread.__init__(self)
         self.handler = handler
         self.host = host
         self.port = port
+        self.requests_to_handle = requests_to_handle
+ 
+        self.wait_to_close_event = wait_to_close_event
         self.ready_event = threading.Event()
         self.stop_event = threading.Event()
 
     @classmethod
-    def basic_response_server(cls, host='localhost', port=0):
-        def basic_response_handler(server_sock):
-            sock, _ = server_sock.accept()
+    def basic_response_server(cls, **kwargs):
+        def basic_response_handler(sock):
             sock.send(
                 b'HTTP/1.1 200 OK\r\n'
                 b'Content-Length: 0\r\n'
                 b'\r\n'
             )
 
-        server = Server(basic_response_handler, host=host, port=port)
+        server = Server(basic_response_handler, **kwargs)
         return server
 
+
     def run(self):
         try:
             sock = self._create_socket_and_bind()
             # in case self.port = 0
             self.port = sock.getsockname()[1]
             self.ready_event.set()
-            self.handler(sock)
-            
+            self._handle_requests_and_close_server(sock)
         finally:
             self.ready_event.set() # just in case of exception
             sock.close()
@@ -45,6 +53,14 @@ class Server(threading.Thread):
         sock.listen(0)
         return sock
 
+    def _handle_requests_and_close_server(self, server_sock):
+        for _ in range(self.requests_to_handle):
+            sock = server_sock.accept()[0]
+            self.handler(sock)
+        
+        if self.wait_to_close_event:
+            self.wait_to_close_event.wait()
+        
     def __enter__(self):
        self.start()
        self.ready_event.wait()