Issue 7019 - [CMS] Refactor `test_server.py`

cms/bin/test_server.py

Index: cms/bin/test_server.py
diff --git a/cms/bin/test_server.py b/cms/bin/test_server.py
old mode 100644
new mode 100755
index 3255cd977b8af1a4fad89f6a324f5487fba2c074..56bc05e96424a2166966309be65097f0e384f1de
--- a/cms/bin/test_server.py
+++ b/cms/bin/test_server.py
@@ -13,20 +13,19 @@
 # You should have received a copy of the GNU General Public License
 # along with Adblock Plus.  If not, see <http://www.gnu.org/licenses/>.
 
-import mimetypes
+from __future__ import print_function
+
 import os
-import sys
+import mimetypes
 import argparse
+import time
 
 import jinja2
 
+from cms.converters import converters
 from cms.utils import process_page
 from cms.sources import create_source
-from cms.converters import converters
-
-source = None
-address = None
-port = None
+from cms.exceptions import ServerSetupTimeoutException
 
 UNICODE_ENCODING = 'utf-8'
 
@@ -47,127 +46,321 @@ ERROR_TEMPLATE = '''
   </body>
 </html>'''
 
-# Initilize the mimetypes modules manually for consistent behavior,
+# Initialize the mimetypes modules manually for consistent behavior,
 # ignoring local files and Windows Registry.
 mimetypes.init([])
 
 
-def get_page(path):
-    path = path.strip('/')
-    if path == '':
-        path = source.read_config().get('general', 'defaultlocale')
-    if '/' in path:
-        locale, page = path.split('/', 1)
-    else:
-        locale, page = path, ''
-
-    default_page = source.read_config().get('general', 'defaultpage')
-    alternative_page = '/'.join([page, default_page]).lstrip('/')
-
-    for format in converters.iterkeys():
-        for p in (page, alternative_page):
-            if source.has_page(p, format):
-                site_url = 'http://{}:{}'.format(address, port)
-                return (p, process_page(source, locale, p, format, site_url))
-    if source.has_localizable_file(locale, page):
-        return (page, source.read_localizable_file(locale, page))
-
-    return (None, None)
-
+class DynamicServerHandler:
+    """General-purpose WSGI server handler that generates pages on request.
 
-def has_conflicting_pages(page):
-    pages = [p for p, _ in source.list_pages()]
-    pages.extend(source.list_localizable_files())
+    Parameters
+    ----------
+    host: str
+        The host where the server will run.
+    port: int
+        The TCP port the server will run on.
+    source_dir: str
+        The path to the website contents.
 
-    if pages.count(page) > 1:
-        return True
-    if any(p.startswith(page + '/') or page.startswith(p + '/') for p in pages):
-        return True
-    return False
+    """
+
+    def __init__(self, host, port, source_dir):
+        self.host = host
+        self.port = port
+        self.source = create_source(source_dir)
+        self.full_url = 'http://{0}:{1}'.format(host, port)
+
+    def _get_data(self, path):
+        """Read the data corresponding to a website path.
+
+        Parameters
+        ----------
+        path: str
+            The path to the page to get the data for.
+
+        Returns
+        -------
+        str or bytes
+            The data corresponding to the path we're trying to access.
+
+        """
+        if self.source.has_static(path):
+            return self.source.read_static(path)
+
+        page, data = self._get_page(path)
+
+        if page and self._has_conflicts(page):
+            raise Exception('The requested page conflicts with another page.')
+
+        return data
+
+    def _get_page(self, path):
+        """Construct a page and return its contents.
+
+        Parameters
+        ----------
+        path: str
+            The path of the page we want to construct.
+
+        Returns
+        -------
+        (page_name, page_contents): (str, str)
+
+        """
+        path = path.strip('/')
+        if path == '':
+            locale, page = self.source.read_config().get('general',
+                                                         'defaultlocale'), ''
+        elif '/' in path:
+            locale, page = path.split('/', 1)
+        else:
+            locale, page = path, ''
+
+        default_page = self.source.read_config().get('general', 'defaultpage')
+        possible_pages = [page, '/'.join([page, default_page]).lstrip('/')]
+
+        for page_format in converters.iterkeys():
+            for p in possible_pages:
+                if self.source.has_page(p, page_format):
+                    return p, process_page(self.source, locale, p, page_format,
+                                           self.full_url)
+
+        if self.source.has_localizable_file(locale, page):
+            return page, self.source.read_localizable_file(locale, page)
+
+        return None, None
+
+    def _has_conflicts(self, page):
+        """Check if a page has conflicts.
+
+        A page has conflicts if there are other pages with the same name.
+
+        Parameters
+        ----------
+        page: str
+            The path of the page we're checking for conflicts.
+
+        Returns
+        -------
+        bool
+            True - if the page has conflicts
+            False - otherwise
+
+        """
+        pages = [p for p, _ in self.source.list_pages()]
+        pages.extend(self.source.list_localizable_files())
+
+        if pages.count(page) > 1:
+            return True
+        if any(p.startswith(page + '/') or page.startswith(p + '/') for p in
+               pages):
+            return True
+        return False
+
+    def get_error_page(self, start_response, status, **kw):
+        """Create and display an error page.
+
+        Parameters
+        ----------
+        start_response: function
+            It will be called before constructing the error page, to setup
+            things like the status of the response and the headers.
+        status: str
+            The status of the response we're sending the error page with.
+            Needs to have the following format:
+                "<status_code> <status_message>"
+        kw: dict
+            Any additional arguments that will be passed onto the `stream`
+            method of a `jinja2 Template`.
+
+        Returns
+        -------
+        generator of utf8 strings
+            Fragments of the corresponding error HTML template.
+
+        """
+        env = jinja2.Environment(autoescape=True)
+        page_template = env.from_string(ERROR_TEMPLATE)
+        mime = 'text/html; encoding={}'.format(UNICODE_ENCODING)
+
+        start_response(status, [('Content-Type', mime)])
+
+        for fragment in page_template.stream(status=status, **kw):
+            yield fragment.encode(UNICODE_ENCODING)
+
+    def __call__(self, environ, start_response):
+        """Execute the handler, according to the WSGI standards.
+
+        Parameters
+        ---------
+        environ: dict
+            The environment under which the page is requested.
+            The requested page must be under the `PATH_INFO` key.
+        start_response: function
+            Used to initiate a response. Must take two arguments, in this
+            order:
+                - Response status, in the format "<code> <message>".
+                - Response headers, as a list of tuples.
+
+        Returns
+        -------
+        list of str
+            With the data for a specific page.
+
+        """
+        path = environ.get('PATH_INFO')
+
+        data = self._get_data(path)
+
+        if data is None:
+            return self.get_error_page(start_response, '404 Not Found',
+                                       uri=path)
+
+        mime = mimetypes.guess_type(path)[0] or 'text/html'
+
+        if isinstance(data, unicode):
+            data = data.encode(UNICODE_ENCODING)
+            mime = '{0}; charset={1}'.format(mime, UNICODE_ENCODING)
+
+        start_response('200 OK', [('Content-Type', mime)])
+        return [data]
+
+
+def parse_arguments():
+    """Set up and parse the arguments required by the script.
+
+    Returns
+    -------
+    argparse.Namespace
+        With the script arguments, as parsed.
+
+    """
+    parser = argparse.ArgumentParser(
+        description='CMS development server created to test pages locally and '
+                    'on-the-fly.',
+    )
+
+    parser.add_argument('path', default=os.curdir, nargs='?',
+                        help='Path to the website we intend to run. If not '
+                             'provided, defaults, to the current directory.')
+    parser.add_argument('--host', default='localhost',
+                        help='Address of the host the server will listen on. '
+                             'Defaults to "localhost".')
+    parser.add_argument('--port', default=5000, type=int,
+                        help='TCP port the server will listen on. Default '
+                             '5000.')
+
+    return parser.parse_args()
+
+
+def wait_for_address(run_fn, handler, timeout=1.0, *args, **kw):
+    """Try to setup the server at a given address.
+
+    If an exception occurs, we sleep for .1 seconds and retry until
+    `timeout` has passed.
+
+    Parameters
+    ----------
+    run_fn: callable
+        Used to initiate the server.
+    handler: DynamicServerHandler
+        Defines the parameters and methods required to handle requests.
+    timeout: float
+        Number of seconds to wait until giving up.
+
+    Raises
+    ------
+    ServerSetupTimeoutException
+        If the server couldn't be setup at the required address after
+        `timeout` passed.
+
+    """
+    start_time = time.time()
 
+    while True:
+        try:
+            with run_fn(handler.host, handler.port, handler, *args, **kw):
+                break
+        except Exception as err:
+            time.sleep(.1)
+            if time.time() - start_time >= timeout:
+                raise ServerSetupTimeoutException(handler.host, handler.port,
+                                                  timeout, err)
+
+
+def run_werkzeug_server(handler, **kw):
+    """Set up a server that uses `werkzeug`.
+
+    Parameters
+    ----------
+    handler: DynamicServerHandler
+        Defines the parameters and methods required to handle requests.
+
+    Raises
+    ------
+    ImportError
+        If the package `werkzeug` is not installed
 
-def get_data(path):
-    if source.has_static(path):
-        return source.read_static(path)
+    """
+    from werkzeug.serving import run_simple
+    import logging
 
-    page, data = get_page(path)
-    if page and has_conflicting_pages(page):
-        raise Exception('The requested page conflicts with another page')
-    return data
+    def run(*args, **kwargs):
+        # The werkzeug logger must be configured before the
+        # root logger. Also we must prevent it from propagating
+        # messages, otherwise messages are logged twice.
+        logger = logging.getLogger('werkzeug')
+        logger.propagate = False
+        logger.setLevel(logging.INFO)
+        logger.addHandler(logging.StreamHandler())
 
+        run_simple(threaded=True, *args, **kwargs)
 
-def show_error(start_response, status, **kwargs):
-    env = jinja2.Environment(autoescape=True)
-    template = env.from_string(ERROR_TEMPLATE)
-    mime = 'text/html; encoding=%s' % UNICODE_ENCODING
-    start_response(status, [('Content-Type', mime)])
-    for fragment in template.stream(status=status, **kwargs):
-        yield fragment.encode(UNICODE_ENCODING)
+    wait_for_address(run, handler, **kw)
 
 
-def handler(environ, start_response):
-    path = environ.get('PATH_INFO')
+def run_builtins_server(handler, **kw):
+    """Configure a server that only uses builtin packages.
+
+    Parameters
+    ----------
+    handler: DynamicServerHandler
+        Defines the parameters and methods required to handle requests.
+
+    """
+    from SocketServer import ThreadingMixIn
+    from wsgiref.simple_server import WSGIServer, make_server
+
+    class ThreadedWSGIServer(ThreadingMixIn, WSGIServer):
+        daemon_threads = True
+
+    def run(host, port, app, **kwargs):
+        def wrapper(environ, start_response):
+            try:
+                return app(environ, start_response)
+            except Exception as e:
+                return handler.get_error_page(
+                    start_response, '500 Internal Server Error',
+                    uri=environ.get('PATH_INFO'), error=e,
+                )
+
+        server = make_server(host, port, wrapper, ThreadedWSGIServer)
+        print(' * Running on {0}:{1}'.format(*server.server_address))
+        server.serve_forever()
 
-    data = get_data(path)
-    if data is None:
-        return show_error(start_response, '404 Not Found', uri=path)
+    wait_for_address(run, handler, **kw)
 
-    mime = mimetypes.guess_type(path)[0] or 'text/html'
 
-    if isinstance(data, unicode):
-        data = data.encode(UNICODE_ENCODING)
-        mime = '%s; charset=%s' % (mime, UNICODE_ENCODING)
-
-    start_response('200 OK', [('Content-Type', mime)])
-    return [data]
-
-
-if __name__ == '__main__':
-
-    parser = argparse.ArgumentParser(description='CMS development server created to test pages locally and on-the-fly')
-    parser.add_argument('path', nargs='?', default=os.curdir)
-    parser.add_argument('-a', '--address', default='localhost', help='Address of the interface the server will listen on')
-    parser.add_argument('-p', '--port', type=int, default=5000, help='TCP port the server will listen on')
-    args = parser.parse_args()
-
-    source = create_source(args.path)
-    address = args.address
-    port = args.port
+def main():
+    args = parse_arguments()
+    handler = DynamicServerHandler(args.host, args.port, args.path)
 
     try:
-        from werkzeug.serving import ThreadedWSGIServer, run_simple
-
-        # see https://github.com/mitsuhiko/werkzeug/pull/770
-        ThreadedWSGIServer.daemon_threads = True
-
-        def run(*args, **kwargs):
-            # The werkzeug logger must be configured before the
-            # root logger. Also we must prevent it from propagating
-            # messages, otherwise messages are logged twice.
-            import logging
-            logger = logging.getLogger('werkzeug')
-            logger.propagate = False
-            logger.setLevel(logging.INFO)
-            logger.addHandler(logging.StreamHandler())
-
-            run_simple(threaded=True, *args, **kwargs)
+        run_werkzeug_server(handler, use_reloader=True, use_debugger=True)
     except ImportError:
-        from SocketServer import ThreadingMixIn
-        from wsgiref.simple_server import WSGIServer, make_server
-
-        class ThreadedWSGIServer(ThreadingMixIn, WSGIServer):
-            daemon_threads = True
+        run_builtins_server(handler)
 
-        def run(host, port, app, **kwargs):
-            def wrapper(environ, start_response):
-                try:
-                    return app(environ, start_response)
-                except Exception as e:
-                    return show_error(start_response, '500 Internal Server Error',
-                                      uri=environ.get('PATH_INFO'), error=e)
 
-            server = make_server(host, port, wrapper, ThreadedWSGIServer)
-            print ' * Running on http://%s:%i/' % server.server_address
-            server.serve_forever()
-
-    run(address, port, handler, use_reloader=True, use_debugger=True)
+if __name__ == '__main__':
+    main()