From e5627134d9a2c052a523f66a2ec9867b3432fae2 Mon Sep 17 00:00:00 2001
From: Matthew Honnibal <honnibal+gh@gmail.com>
Date: Sun, 23 Oct 2016 15:50:04 +0200
Subject: [PATCH] Test Issue #461: ent_iob tag incorrect after setting
 entities.

---
 spacy/tests/tokens/test_add_entities.py | 42 +++++++++++++++++++++++++
 1 file changed, 42 insertions(+)
 create mode 100644 spacy/tests/tokens/test_add_entities.py

diff --git a/spacy/tests/tokens/test_add_entities.py b/spacy/tests/tokens/test_add_entities.py
new file mode 100644
index 000000000..8d789714f
--- /dev/null
+++ b/spacy/tests/tokens/test_add_entities.py
@@ -0,0 +1,42 @@
+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
+
+import pytest
+
+
+@pytest.fixture(scope="module")
+def en_vocab():
+    return spacy.get_lang_class('en').Defaults.create_vocab()
+
+
+@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))
+    doc.ents = [(doc.vocab.strings['ANIMAL'], 3, 4)]
+    tags = [w.ent_iob_ for w in doc]
+    assert tags == ['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']