Handle relationships in notes

1 files changed, 9 insertions, 7 deletions

diff --git a/tests/docx/notes_xml_tests.py b/tests/docx/notes_xml_tests.py
index 90de8f8..b912991 100644
--- a/tests/docx/notes_xml_tests.py
+++ b/tests/docx/notes_xml_tests.py
@@ -1,18 +1,20 @@
 from nose.tools import istest, assert_equal
 
+from mammoth import documents
 from mammoth.docx.xmlparser import element as xml_element
-from mammoth.docx.notes_xml import read_footnotes_xml_element
+from mammoth.docx.notes_xml import create_footnotes_reader
+from mammoth.docx import body_xml
 
 
 @istest
 def id_and_body_of_footnote_are_read():
     footnote_body = [xml_element("w:p")]
-    footnotes = read_footnotes_xml_element(xml_element("w:footnotes", {}, [
+    footnotes = create_footnotes_reader(body_xml.reader())(xml_element("w:footnotes", {}, [
         xml_element("w:footnote", {"w:id": "1"}, footnote_body),
     ]))
-    assert_equal(1, len(footnotes))
-    assert footnote_body is footnotes[0].body
-    assert_equal("1", footnotes[0].id)
+    assert_equal(1, len(footnotes.value))
+    assert isinstance(footnotes.value[0].body[0], documents.Paragraph)
+    assert_equal("1", footnotes.value[0].note_id)
 
 
 @istest
@@ -27,8 +29,8 @@ def separator_is_ignored():
 
 def _assert_footnote_type_is_ignored(footnote_type):
     footnote_body = [xml_element("w:p")]
-    footnotes = read_footnotes_xml_element(xml_element("w:footnotes", {}, [
+    footnotes = create_footnotes_reader(None)(xml_element("w:footnotes", {}, [
         xml_element("w:footnote", {"w:id": "1", "w:type": footnote_type}, footnote_body),
     ]))
-    assert_equal(0, len(footnotes))
+    assert_equal(0, len(footnotes.value))