1 files changed, 16 insertions, 10 deletions
diff --git a/python/tests/test_grobid2json.py b/python/tests/test_grobid2json.py
index 8497b10..b00a88d 100644
--- a/python/tests/test_grobid2json.py
+++ b/python/tests/test_grobid2json.py
@@ -1,22 +1,28 @@
-
-import xml
 import json
+import xml
+
 import pytest
-from grobid2json import *
+from grobid_tei_xml import parse_document_xml
 
 
 def test_small_xml():
-    
-    with open('tests/files/small.xml', 'r') as f:
+    """
+    This used to be a test of grobid2json; now it is a compatability test for
+    the to_legacy_dict() feature of grobid_tei_xml.
+    """
+
+    with open("tests/files/small.xml", "r") as f:
         tei_xml = f.read()
-    with open('tests/files/small.json', 'r') as f:
-        json_form  = json.loads(f.read())
+    with open("tests/files/small.json", "r") as f:
+        json_form = json.loads(f.read())
+
+    tei_doc = parse_document_xml(tei_xml)
+    assert tei_doc.to_legacy_dict() == json_form
 
-    assert teixml2json(tei_xml) == json_form
 
 def test_invalid_xml():
 
     with pytest.raises(xml.etree.ElementTree.ParseError):
-        teixml2json("this is not XML")
+        parse_document_xml("this is not XML")
     with pytest.raises(ValueError):
-        teixml2json("<xml></xml>")
+        parse_document_xml("<xml></xml>")