client/tests/httpserver.py

# Copyright 2014 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.

import json
import logging
import threading

from six.moves import BaseHTTPServer
from six.moves import http_client

_STOP_EVENT = '/fakeserver/__stop__'


class Handler(BaseHTTPServer.BaseHTTPRequestHandler):
  """Handlers implements utility functions to help implementing a fake."""

  ### Public methods

  def send_json(self, data):
    """Sends a JSON response."""
    self.send_response(200)
    self.send_header('Content-type', 'application/json')
    self.end_headers()
    self.wfile.write(json.dumps(data).encode())

  def send_octet_stream(self, data, headers=None):
    """Sends a binary response."""
    self.send_response(200)
    self.send_header('Content-type', 'application/octet-stream')
    for key, value in (headers or {}).items():
      self.send_header(key, value)
    self.end_headers()
    self.wfile.write(data)

  def read_body(self):
    """Reads the request body."""
    return self.rfile.read(int(self.headers['Content-Length']))

  def yield_body(self):
    """Yields the request body as 4kiB chunks."""
    size = int(self.headers['Content-Length'])
    while size:
      chunk = min(4096, size)
      yield self.rfile.read(chunk)
      size -= chunk

  ### Overrides from BaseHTTPRequestHandler

  def do_OPTIONS(self):
    if self.path == _STOP_EVENT:
      self.server.parent._stopped = True
    self.send_octet_stream(b'')

  def log_message(self, fmt, *args):
    logging.info(
        '%s - - [%s] %s', self.address_string(), self.log_date_time_string(),
        fmt % args)


class Server:
  """Server implements a simple HTTP server to implement a fake."""
  _HANDLER_CLS = None

  def __init__(self):
    assert issubclass(self._HANDLER_CLS, Handler), self._HANDLER_CLS
    self._closed = False
    self._stopped = False
    self._server = BaseHTTPServer.HTTPServer(
        ('127.0.0.1', 0), self._HANDLER_CLS)
    self._server.parent = self
    self._server.url = self.url = 'http://127.0.0.1:%d' % (
        self._server.server_port)
    self._thread = threading.Thread(target=self._run, name='httpd')
    self._thread.daemon = True
    self._thread.start()
    logging.info('%s', self.url)

  def close(self):
    assert not self._closed
    self._closed = True
    self._send_event(_STOP_EVENT)
    self._thread.join()

  def _run(self):
    while not self._stopped:
      self._server.handle_request()
    self._server.server_close()

  def _send_event(self, path):
    conn = http_client.HTTPConnection(
        '127.0.0.1:%d' % self._server.server_port, timeout=60)
    try:
      conn.request('OPTIONS', path)
      conn.getresponse()
    finally:
      conn.close()