From 05447be8846655d1803884aafadc103cd46243a4 Mon Sep 17 00:00:00 2001
From: Ines Montani <ines@ines.io>
Date: Wed, 11 Jan 2017 18:54:24 +0100
Subject: [PATCH] Modernise test for adding entities

---
 spacy/tests/tokens/test_add_entities.py | 47 ++++++++-----------------
 1 file changed, 14 insertions(+), 33 deletions(-)

diff --git a/spacy/tests/tokens/test_add_entities.py b/spacy/tests/tokens/test_add_entities.py
index 8d789714f..e4d62398d 100644
--- a/spacy/tests/tokens/test_add_entities.py
+++ b/spacy/tests/tokens/test_add_entities.py
@@ -1,42 +1,23 @@
+# coding: utf-8
 from __future__ import unicode_literals
-import spacy
-from spacy.vocab import Vocab
-from spacy.matcher import Matcher
-from spacy.tokens.doc import Doc
-from spacy.attrs import *
-from spacy.pipeline import EntityRecognizer
+
+from ...pipeline import EntityRecognizer
+from ..util import get_doc
 
 import pytest
 
 
-@pytest.fixture(scope="module")
-def en_vocab():
-    return spacy.get_lang_class('en').Defaults.create_vocab()
+def test_add_entities_set_ents_iob(en_vocab):
+    text = ["This", "is", "a", "lion"]
+    doc = get_doc(en_vocab, text)
+    ner = EntityRecognizer(en_vocab, features=[(2,), (3,)])
+    ner(doc)
 
-
-@pytest.fixture(scope="module")
-def entity_recognizer(en_vocab):
-    return EntityRecognizer(en_vocab, features=[(2,), (3,)])
-
-@pytest.fixture
-def animal(en_vocab):
-    return nlp.vocab.strings[u"ANIMAL"]
-
-
-@pytest.fixture
-def doc(en_vocab, entity_recognizer):
-    doc = Doc(en_vocab, words=[u"this",  u"is", u"a", u"lion"])
-    entity_recognizer(doc)
-    return doc
-
-
-def test_set_ents_iob(doc):
     assert len(list(doc.ents)) == 0
-    tags = [w.ent_iob_ for w in doc]
-    assert tags == (['O'] * len(doc))
+    assert [w.ent_iob_ for w in doc] == (['O'] * len(doc))
+
     doc.ents = [(doc.vocab.strings['ANIMAL'], 3, 4)]
-    tags = [w.ent_iob_ for w in doc]
-    assert tags == ['O', 'O', 'O', 'B']
+    assert [w.ent_iob_ for w in doc] == ['O', 'O', 'O', 'B']
+
     doc.ents = [(doc.vocab.strings['WORD'], 0, 2)]
-    tags = [w.ent_iob_ for w in doc]
-    assert tags == ['B', 'I', 'O', 'O']
+    assert [w.ent_iob_ for w in doc] == ['B', 'I', 'O', 'O']