import socket
import http.server
import threading
import unittest
from time import sleep

RESPONSE_TEXT = b'Expected response'
PORT = 8123

class DummyRequestHandler(http.server.BaseHTTPRequestHandler):
    def log_message(self, *args):
        # no log messages
        pass
        
    def do_GET(self):
        self.send_response(200)
        self.send_header('Content-Type', 'text/plain')
        self.send_header('Content-Length', len(RESPONSE_TEXT))
        self.send_header('Connection', 'close')
        self.end_headers()
        self.wfile.write(RESPONSE_TEXT)

class DummyRequestHandlerWithCycle(DummyRequestHandler):
    def do_GET(self):
        # this cycle makes the server not to close the connection
        self.foo = self
        DummyRequestHandler.do_GET(self)
        
class TestHttpClose(unittest.TestCase):

  def run_server_bg(self, handler_class):
    httpd = http.server.HTTPServer(('127.0.0.1', PORT), handler_class)
    httpd.handle_request() # only 1 request!

  def _test(self, handler_class):
    t = threading.Thread(target=self.run_server_bg, args=(handler_class,))
    t.start()
    sleep(1)
    s = socket.socket()
    s.settimeout(5)
    s.connect(('127.0.0.1', PORT))
    req = b'GET / HTTP/1.0\r\n\r\n'
    s.sendall(req)
    resp = b''
    while True:
      d = s.recv(100)
      if not d: break
      resp += d
    self.assert_(RESPONSE_TEXT in resp)
    self.assert_(s.recv(1)==b'', "Connection should be closed by now")

  def testCloseNoCycle(self):
    self._test(handler_class=DummyRequestHandler)
    
  def testCloseWithCycle(self):
    self._test(handler_class=DummyRequestHandlerWithCycle)
    
if __name__=='__main__':
   unittest.main()