def test_write_header(self, write_method):
writer = HttpWriter(None, None, None, None)
writer.status = 200
writer['foo'] = 'bar'
writer.flush()
write_method.assert_called_with(b'HTTP/1.1 200 OK\r\nfoo: bar\r\n\r\n')
def test_write_status(self, write_method):
writer = HttpWriter(None, None, None, None)
writer.write_status(b'200 OK')
self.assertFalse(writer._headers_sent)
writer.flush()
self.assertTrue(writer._headers_sent)
write_method.assert_called_with(b'HTTP/1.1 200 OK\r\n\r\n')
def test_status_written(self):
writer = HttpWriter(None, None, None, None)
self.assertFalse(writer._headers_sent)
writer._status_written = self._headers_sent = True
writer.restore()
self.assertFalse(writer._headers_sent)
def test_write_header(self, write_method):
writer = HttpWriter(None, None, None, None)
writer.write_status(b'200 OK')
writer.write_header(b'foo', b'bar')
writer.flush()
write_method.assert_called_with(b'HTTP/1.1 200 OK\r\nfoo: bar\r\n\r\n')
def test_maybe_finalize_headers(self, write_method):
writer = HttpWriter(None, None, None, None)
writer._maybe_send_headers()
write_method.assert_called_with(b'\r\n')
writer = HttpWriter(None, None, None, None)
writer._headers_sent = True
writer._maybe_send_headers()
self.assertTrue(writer._headers_sent)
def connection_made(self, transport):
self._transport = transport
self._reader.set_transport(transport)
self._writer = HttpWriter(transport, self,
self._reader,
self._loop)
self._handler = self._build_handler()
self._task = asyncio.async(self._handle_client())
self._task.add_done_callback(self._maybe_log_exception)
self._reset_timeout()
def test_writelines(self):
mtransport = unittest.mock.MagicMock()
writer = HttpWriter(mtransport, None, None, None)
writer.writelines((b'Hello',))
self.assertTrue(writer._headers_sent)
mtransport.writelines.assert_called_with((b'Hello',))
def test_write_body(self, write_method):
writer = HttpWriter(None, None, None, None)
writer.write_body(b'hello')
self.assertTrue(writer._headers_sent)
write_method.assert_called_with(b'hello')
def test_write_headers(self, write_header_method):
writer = HttpWriter(None, None, None, None)
writer.write_headers(((b'foo', b'bar'),))
write_header_method.assert_called_with(b'foo', b'bar')
def test_write_header_raises_when_headers_sent(self):
writer = HttpWriter(None, None, None, None)
writer._headers_sent = True
self.assertRaises(AssertionError, writer.write_header, b'foo', b'bar')
def test_write_headers(self, write_header_method):
writer = HttpWriter(None, None, None, None)
writer.add_headers(('foo', 'bar'))
write_header_method.assert_called_with('foo', 'bar')
class BaseHttpProtocol(asyncio.StreamReaderProtocol):
handler_factory = BaseHandler
def __init__(self, *, keep_alive=None, loop=None):
if keep_alive is None:
keep_alive = _DEFAULT_KEEP_ALIVE
self._reader = asyncio.StreamReader(loop=loop)
self._keep_alive = keep_alive
super().__init__(self._reader, None, loop)
self.h_timeout = None
def connection_made(self, transport):
self._transport = transport
self._reader.set_transport(transport)
self._writer = HttpWriter(transport, self,
self._reader,
self._loop)
self._handler = self._build_handler()
self._task = asyncio.async(self._handle_client())
self._task.add_done_callback(self._maybe_log_exception)
self._reset_timeout()
def connection_lost(self, exc):
self._task.cancel()
self._task = None
handler = self._handler
self._writer = None
self._handler = None
self._stop_timeout()
try:
handler.connection_lost(exc)
finally:
super().connection_lost(exc)
def data_received(self, data):
self._reset_timeout()
super().data_received(data)
@asyncio.coroutine
def _handle_client(self):
while True:
try:
req = yield from HttpParser.parse(self._reader)
except BadRequestException as e:
self._bad_request(self._writer, e)
self._writer.close()
break
# connection has been closed
if req is None:
break
try:
yield from self._handler.handle_request(req)
finally:
if self._should_close_conn_immediately(req):
self._writer.close()
else:
yield from req.body.read()
if self._writer is not None:
self._writer.restore()
def _reset_timeout(self):
self._stop_timeout()
self.h_timeout = self._loop.call_later(
self._keep_alive, self._handle_timeout)
def _stop_timeout(self):
if self.h_timeout is not None:
self.h_timeout.cancel()
self.h_timeout = None
def _handle_timeout(self):
if self._handler.on_timeout():
self._reset_timeout()
def _build_handler(self):
return self.handler_factory(self._transport, self,
self._reader,
self._writer,
**self.get_handler_kwargs())
def get_handler_kwargs(self):
return {}
def _should_close_conn_immediately(self, req):
if self._keep_alive < 1:
return True
should_close = False
if req.version.lower() == 'http/1.0':
should_close = True
conn_header = req['connection']
if not conn_header:
return should_close
if conn_header == 'keep-alive':
should_close = False
elif conn_header == 'close':
should_close = True
return should_close
def _maybe_log_exception(self, task):
try:
if not task.cancelled():
task.result()
except:
logger.exception("An exception ocurred while serving request")
if self._writer is not None:
self._writer.close()