flaskapp / werkzeug / contrib /

# -*- coding: utf-8 -*-

    This module implements extended wrappers for simplified testing.

        A response wrapper which adds various cached attributes for
        simplified assertions on various content types.

    :copyright: (c) 2010 by the Werkzeug Team, see AUTHORS for more details.
    :license: BSD, see LICENSE for more details.
from werkzeug import Response, cached_property, import_string

class ContentAccessors(object):
    A mixin class for response objects that provides a couple of useful
    accessors for unittesting.

    def xml(self):
        """Get an etree if possible."""
        if 'xml' not in self.mimetype:
            raise AttributeError(
                'Not a XML response (Content-Type: %s)'
                % self.mimetype)
        for module in ['xml.etree.ElementTree', 'ElementTree',
            etree = import_string(module, silent=True)
            if etree is not None:
                return etree.XML(self.body)
        raise RuntimeError('You must have ElementTree installed '
                           'to use TestResponse.xml')
    xml = cached_property(xml)

    def lxml(self):
        """Get an lxml etree if possible."""
        if ('html' not in self.mimetype and 'xml' not in self.mimetype):
            raise AttributeError('Not an HTML/XML response')
        from lxml import etree
            from lxml.html import fromstring
        except ImportError:
            fromstring = etree.HTML
        if self.mimetype=='text/html':
            return fromstring(
        return etree.XML(
    lxml = cached_property(lxml)

    def json(self):
        """Get the result of simplejson.loads if possible."""
        if 'json' not in self.mimetype:
            raise AttributeError('Not a JSON response')
            from simplejson import loads
            from json import loads
        return loads(
    json = cached_property(json)

class TestResponse(Response, ContentAccessors):
    """Pass this to `werkzeug.test.Client` for easier unittesting."""