1 files changed, 54 insertions, 5 deletions
diff --git a/Lib/test/test_pyexpat.py b/Lib/test/test_pyexpat.py
index 5bb8c97..27eecb8 100644
--- a/Lib/test/test_pyexpat.py
+++ b/Lib/test/test_pyexpat.py
@@ -155,6 +155,14 @@ class ParseTest(unittest.TestCase):
         'ElementDeclHandler', 'AttlistDeclHandler', 'SkippedEntityHandler',
         ]
 
+    def _hookup_callbacks(self, parser, handler):
+        """
+        Set each of the callbacks defined on handler and named in
+        self.handler_names on the given parser.
+        """
+        for name in self.handler_names:
+            setattr(parser, name, getattr(handler, name))
+
     def _verify_parse_output(self, operations):
         expected_operations = [
             ('XML declaration', ('1.0', 'iso-8859-1', 0)),
@@ -196,8 +204,7 @@ class ParseTest(unittest.TestCase):
         # Try the parse again, this time producing Unicode output
         out = self.Outputter()
         parser = expat.ParserCreate(namespace_separator='!')
-        for name in self.handler_names:
-            setattr(parser, name, getattr(out, name))
+        self._hookup_callbacks(parser, out)
 
         parser.Parse(data, 1)
 
@@ -210,8 +217,7 @@ class ParseTest(unittest.TestCase):
         # Try parsing a file
         out = self.Outputter()
         parser = expat.ParserCreate(namespace_separator='!')
-        for name in self.handler_names:
-            setattr(parser, name, getattr(out, name))
+        self._hookup_callbacks(parser, out)
         file = BytesIO(data)
 
         parser.ParseFile(file)
@@ -613,6 +619,48 @@ class ErrorMessageTest(unittest.TestCase):
                              errors.codes[errors.XML_ERROR_UNCLOSED_TOKEN])
 
 
+class ForeignDTDTests(unittest.TestCase):
+    """
+    Tests for the UseForeignDTD method of expat parser objects.
+    """
+    def test_use_foreign_dtd(self):
+        """
+        If UseForeignDTD is passed True and a document without an external
+        entity reference is parsed, ExternalEntityRefHandler is first called
+        with None for the public and system ids.
+        """
+        handler_call_args = []
+        def resolve_entity(context, base, system_id, public_id):
+            handler_call_args.append((public_id, system_id))
+            return 1
+
+        parser = expat.ParserCreate()
+        parser.UseForeignDTD(True)
+        parser.SetParamEntityParsing(expat.XML_PARAM_ENTITY_PARSING_ALWAYS)
+        parser.ExternalEntityRefHandler = resolve_entity
+        parser.Parse("<?xml version='1.0'?><element/>")
+        self.assertEqual(handler_call_args, [(None, None)])
+
+    def test_ignore_use_foreign_dtd(self):
+        """
+        If UseForeignDTD is passed True and a document with an external
+        entity reference is parsed, ExternalEntityRefHandler is called with
+        the public and system ids from the document.
+        """
+        handler_call_args = []
+        def resolve_entity(context, base, system_id, public_id):
+            handler_call_args.append((public_id, system_id))
+            return 1
+
+        parser = expat.ParserCreate()
+        parser.UseForeignDTD(True)
+        parser.SetParamEntityParsing(expat.XML_PARAM_ENTITY_PARSING_ALWAYS)
+        parser.ExternalEntityRefHandler = resolve_entity
+        parser.Parse(
+            "<?xml version='1.0'?><!DOCTYPE foo PUBLIC 'bar' 'baz'><element/>")
+        self.assertEqual(handler_call_args, [("bar", "baz")])
+
+
 def test_main():
     run_unittest(SetAttributeTest,
                  ParseTest,
@@ -624,7 +672,8 @@ def test_main():
                  sf1296433Test,
                  ChardataBufferTest,
                  MalformedInputTest,
-                 ErrorMessageTest)
+                 ErrorMessageTest,
+                 ForeignDTDTests)
 
 if __name__ == "__main__":
     test_main()