import functools


def content_type(func, content_type='json'):
    """
    Decorator for a test case method that sets the test case's
    content_type to 'json' or 'xml' and resets it afterwards to
    the original setting. This also asserts that if there is a
    value for the test object's `res` attribute, that the content-type
    header of the response is correct.
    """
    @functools.wraps(func)
    def wrapped(*a, **kwargs):
        test_obj = a[0]
        orig_content_type = test_obj.content_type
        try:
            test_obj.content_type = content_type
            func(*a, **kwargs)
            if getattr(test_obj, 'res'):
                expected = 'application/%s' % content_type
                got = test_obj.res.headers['content-type'].split(';')[0]
                test_obj.assertEqual(expected, got,
                                     "Bad content type: %s. Expected: %s" %
                                     (got, expected))
        finally:
            test_obj.content_type = orig_content_type
    return wrapped