added test code for error cases

1 files changed, 92 insertions, 15 deletions

diff --git a/tests.py b/tests.py
index 23a9802..8e55e1e 100644
--- a/tests.py
+++ b/tests.py
@@ -8,19 +8,21 @@ from StringIO import StringIO
 from xml.sax.saxutils import XMLGenerator
 
 from defusedxml import cElementTree, ElementTree, minidom, pulldom, sax
+from defusedxml import DTDForbidden, EntityForbidden, NotSupportedError
+
 
 try:
     from defusedxml import lxml
+    LXML3 = lxml.LXML3
 except ImportError:
     lxml = None
+    LXML3 = False
 
 
 HERE = os.path.dirname(os.path.abspath(__file__))
 
 class BaseTests(unittest.TestCase):
     module = None
-    parse = None
-    parseString = None
 
     xml_dtd = os.path.join(HERE, "xmltestdata", "dtd.xml")
     xml_external = os.path.join(HERE, "xmltestdata", "external.xml")
@@ -30,11 +32,16 @@ class BaseTests(unittest.TestCase):
     xml_bomb = os.path.join(HERE, "xmltestdata", "xmlbomb.xml")
 
     def setUp(self):
-        self.parse = self.module.parse
-        if hasattr(self.module, "fromstring"):
-            self.parseString = self.module.fromstring
-        else:
-            self.parseString = self.module.parseString
+        if not hasattr(self, "parse"):
+            self.parse = self.module.parse
+        if not hasattr(self, "parseString"):
+            if hasattr(self.module, "fromstring"):
+                self.parseString = self.module.fromstring
+            else:
+                self.parseString = self.module.parseString
+        if not hasattr(self, "iterparse"):
+            if hasattr(self.module, "iterparse"):
+                self.iterparse = self.module.iterparse
 
     def get_content(self, xmlfile):
         with io.open(xmlfile, "rb") as f:
@@ -43,6 +50,50 @@ class BaseTests(unittest.TestCase):
     def test_simple_parse(self):
         self.parse(self.xml_simple)
         self.parseString(self.get_content(self.xml_simple))
+        if self.iterparse:
+            self.iterparse(self.xml_simple)
+
+    def test_simple_parse_ns(self):
+        self.parse(self.xml_simple_ns)
+        self.parseString(self.get_content(self.xml_simple_ns))
+        if self.iterparse:
+            self.iterparse(self.xml_simple_ns)
+
+    def test_entities_forbidden(self):
+        self.assertRaises(EntityForbidden, self.parse, self.xml_bomb)
+        self.assertRaises(EntityForbidden, self.parse, self.xml_quadratic)
+        self.assertRaises(EntityForbidden, self.parse, self.xml_external)
+
+        #self.parse(self.xml_dtd)
+        self.assertRaises(EntityForbidden, self.parseString,
+                          self.get_content(self.xml_bomb))
+        self.assertRaises(EntityForbidden, self.parseString,
+                          self.get_content(self.xml_quadratic))
+        self.assertRaises(EntityForbidden, self.parseString,
+                          self.get_content(self.xml_external))
+
+    def test_dtd_forbidden(self):
+        self.assertRaises(DTDForbidden, self.parse, self.xml_bomb,
+                          forbid_dtd=True)
+        self.assertRaises(DTDForbidden, self.parse, self.xml_quadratic,
+                          forbid_dtd=True)
+        self.assertRaises(DTDForbidden, self.parse, self.xml_external,
+                          forbid_dtd=True)
+        self.assertRaises(DTDForbidden, self.parse, self.xml_dtd,
+                          forbid_dtd=True)
+
+        self.assertRaises(DTDForbidden, self.parseString,
+                          self.get_content(self.xml_bomb),
+                          forbid_dtd=True)
+        self.assertRaises(DTDForbidden, self.parseString,
+                          self.get_content(self.xml_quadratic),
+                          forbid_dtd=True)
+        self.assertRaises(DTDForbidden, self.parseString,
+                          self.get_content(self.xml_external),
+                          forbid_dtd=True)
+        self.assertRaises(DTDForbidden, self.parseString,
+                          self.get_content(self.xml_dtd),
+                          forbid_dtd=True)
 
 
 class TestDefusedcElementTree(BaseTests):
@@ -56,34 +107,60 @@ class TestDefusedElementTree(BaseTests):
 class TestDefusedMinidom(BaseTests):
     module = minidom
 
+    iterparse = None
+
 
 class TestDefusedPulldom(BaseTests):
     module = pulldom
 
+    iterparse = None
+
+    def parse(self, xmlfile, **kwargs):
+        dom = self.module.parse(xmlfile, **kwargs)
+        list(dom)
+
+    def parseString(self, xmlstring, **kwargs):
+        dom = self.module.parseString(xmlstring, **kwargs)
+        list(dom)
 
 class TestDefusedSax(BaseTests):
     module = sax
 
-    def test_simple_parse(self):
+    iterparse = None
+
+    def parse(self, xmlfile, **kwargs):
+        result = StringIO()
+        handler = XMLGenerator(result)
+        self.module.parse(xmlfile, handler, **kwargs)
+
+    def parseString(self, xmlstring, **kwargs):
         result = StringIO()
-        gen = XMLGenerator(result)
-        self.module.parse(self.xml_simple, gen)
+        handler = XMLGenerator(result)
+        self.module.parseString(xmlstring, handler, **kwargs)
 
 
 class TestDefusedLxml(BaseTests):
     module = lxml
 
+    iterparse = None
+
+    if not LXML3:
+        def test_entities_forbidden(self):
+            self.assertRaises(NotSupportedError, self.parse, self.xml_bomb)
+
 
 def test_main():
     suite = unittest.TestSuite()
-    for cls in (TestDefusedcElementTree, TestDefusedElementTree,
-                TestDefusedMinidom, TestDefusedPulldom,
-                TestDefusedSax):
-        suite.addTests(unittest.makeSuite(cls))
+    suite.addTests(unittest.makeSuite(TestDefusedcElementTree))
+    suite.addTests(unittest.makeSuite(TestDefusedElementTree))
+    suite.addTests(unittest.makeSuite(TestDefusedMinidom))
+    suite.addTests(unittest.makeSuite(TestDefusedPulldom))
+    suite.addTests(unittest.makeSuite(TestDefusedSax))
     if lxml is not None:
         suite.addTests(unittest.makeSuite(TestDefusedLxml))
     return suite
 
 if __name__ == "__main__":
-    result = unittest.TextTestRunner(verbosity=2).run(test_main())
+    suite = test_main()
+    result = unittest.TextTestRunner(verbosity=2).run(suite)
     sys.exit(not result.wasSuccessful())