Update test_displacy.py
@@ -50,20 +50,6 @@ def test_displacy_spans(en_vocab):
assert html.startswith("<div")
-def test_displacy_render_wrapper(en_vocab):
- """Test that displaCy accepts custom rendering wrapper."""
-
- def wrapper(html):
- return "TEST" + html + "TEST"
-
- displacy.set_render_wrapper(wrapper)
- doc = get_doc(en_vocab, words=["But", "Google", "is", "starting", "from", "behind"])
- doc.ents = [Span(doc, 1, 2, label=doc.vocab.strings["ORG"])]
- html = displacy.render(doc, style="ent")
- assert html.startswith("TEST<div")
- assert html.endswith("/div>TEST")
-
-
def test_displacy_raises_for_wrong_type(en_vocab):
with pytest.raises(ValueError):
displacy.render("hello world")
@@ -86,3 +72,18 @@ def test_displacy_rtl():
html = displacy.render(doc, page=True, style="ent")
assert "direction: rtl" in html
assert 'lang="{}"'.format(nlp.lang) in html
+
+def test_displacy_render_wrapper(en_vocab):
+ """Test that displaCy accepts custom rendering wrapper."""
+
+ def wrapper(html):
+ return "TEST" + html + "TEST"
+
+ displacy.set_render_wrapper(wrapper)
+ doc = get_doc(en_vocab, words=["But", "Google", "is", "starting", "from", "behind"])
+ doc.ents = [Span(doc, 1, 2, label=doc.vocab.strings["ORG"])]
+ html = displacy.render(doc, style="ent")
+ assert html.startswith("TEST<div")
+ assert html.endswith("/div>TEST")
+ # Restore
+ displacy.set_render_wrapper(lambda html: html)
\ No newline at end of file
