Restore tests for mention_extractor.get_mentions(docs) and candidate_…

…extractor.get_candidates(docs)
HazyResearch · Jun 4, 2020 · c30139f · c30139f
1 parent 02f35e4
commit c30139f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/e2e/test_e2e.py b/tests/e2e/test_e2e.py
@@ -152,6 +152,7 @@ def test_e2e():
     mention_extractor.apply(docs, parallelism=PARALLEL)
 
     assert len(mention_extractor.get_mentions()) == 3
+    assert len(mention_extractor.get_mentions(docs)) == 3
 
     # Candidate Extraction
     candidate_extractor = CandidateExtractor(
@@ -177,6 +178,7 @@ def test_e2e():
     )
 
     assert len(train_cands) == 2
+    assert len(candidate_extractor.get_candidates(docs)) == 2
 
     # Featurization
     featurizer = Featurizer(session, [PartTemp, PartVolt])