bundled/flask/flask/testing.py @ f33efe14bff1

# -*- coding: utf-8 -*-
"""
    flask.testing
    ~~~~~~~~~~~~~

    Implements test support helpers.  This module is lazily imported
    and usually not used in production environments.

    :copyright: (c) 2010 by Armin Ronacher.
    :license: BSD, see LICENSE for more details.
"""

from werkzeug import Client, EnvironBuilder
from flask import _request_ctx_stack


class FlaskClient(Client):
    """Works like a regular Werkzeug test client but has some
    knowledge about how Flask works to defer the cleanup of the
    request context stack to the end of a with body when used
    in a with statement.
    """

    preserve_context = context_preserved = False

    def open(self, *args, **kwargs):
        if self.context_preserved:
            _request_ctx_stack.pop()
            self.context_preserved = False
        kwargs.setdefault('environ_overrides', {}) \
            ['flask._preserve_context'] = self.preserve_context

        as_tuple = kwargs.pop('as_tuple', False)
        buffered = kwargs.pop('buffered', False)
        follow_redirects = kwargs.pop('follow_redirects', False)

        builder = EnvironBuilder(*args, **kwargs)

        if self.application.config.get('SERVER_NAME'):
            server_name = self.application.config.get('SERVER_NAME')
            if ':' not in server_name:
                http_host, http_port = server_name, None
            else:
                http_host, http_port = server_name.split(':', 1)
            if builder.base_url == 'http://localhost/':
                # Default Generated Base URL
                if http_port != None:
                    builder.host = http_host + ':' + http_port
                else:
                    builder.host = http_host
        old = _request_ctx_stack.top
        try:
            return Client.open(self, builder,
                               as_tuple=as_tuple,
                               buffered=buffered,
                               follow_redirects=follow_redirects)
        finally:
            self.context_preserved = _request_ctx_stack.top is not old

    def __enter__(self):
        self.preserve_context = True
        return self

    def __exit__(self, exc_type, exc_value, tb):
        self.preserve_context = False
        if self.context_preserved:
            _request_ctx_stack.pop()