Fixed #236: app.lxml support xml responses with encoding declaration

2 files changed, 30 insertions, 5 deletions

diff --git a/tests/test_response.py b/tests/test_response.py
index ade096c..f2799bb 100644
--- a/tests/test_response.py
+++ b/tests/test_response.py
@@ -1,6 +1,3 @@
-import sys
-
-
 import webtest
 from webtest.debugapp import debug_app
 from webob import Request
@@ -117,6 +114,17 @@ def links_app(environ, start_response):
     return [body]
 
 
+def svg_application(env, start_response):
+    start_response('200 OK', [('Content-Type', 'image/svg+xml')])
+    return [
+        b"""<?xml version="1.0" encoding="UTF-8" standalone="no"?>"""
+        b"""<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN" """
+        b""" "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">"""
+        b"""<svg xmlns="http://www.w3.org/2000/svg" """
+        b"""xmlns:xlink="http://www.w3.org/1999/xlink"></svg>"""
+    ]
+
+
 def gzipped_app(environ, start_response):
     status = "200 OK"
     encoded_body = list(gzip_app_iter([b'test']))
@@ -279,7 +287,6 @@ class TestResponse(unittest.TestCase):
         resp.content_type = 'text/xml'
         resp.xml
 
-    @unittest.skipIf('PyPy' in sys.version, 'skip lxml tests on pypy')
     def test_lxml_attribute(self):
         app = webtest.TestApp(links_app)
         resp = app.post('/')
@@ -287,6 +294,12 @@ class TestResponse(unittest.TestCase):
         print(resp.body)
         print(resp.lxml)
 
+    def test_lxml_attribute_with_encoding_declaration(self):
+        app = webtest.TestApp(svg_application)
+        resp = app.get('/')
+        print(resp.body)
+        print(resp.lxml)
+
     def test_html_attribute(self):
         app = webtest.TestApp(links_app)
         res = app.post('/')
diff --git a/webtest/response.py b/webtest/response.py
index 76917c7..8c4f87e 100644
--- a/webtest/response.py
+++ b/webtest/response.py
@@ -455,7 +455,19 @@ class TestResponse(webob.Response):
         if self.content_type == 'text/html':
             return fromstring(self.testbody, base_url=self.request.url)
         else:
-            return etree.XML(self.testbody, base_url=self.request.url)
+            try:
+                return etree.XML(self.testbody, base_url=self.request.url)
+            except ValueError as e:
+                # Encoding may be declared in xml. Check the error
+                error = ' '.join(e.args)
+                expected_error = (
+                    'Unicode strings with encoding declaration '
+                    'are not supported.'
+                )
+                if error.startswith(expected_error):
+                    # Use bytes
+                    return etree.XML(self.body, base_url=self.request.url)
+                raise
 
     @property
     def json(self):