haystack/test/test_finder.py
Lalit Pagaria abda994116
Pytest fix memory leak and put pytest marker on slow tests (#520)
* Clear faiss_index during teardown

* Marking slow test with pytest markers. So In future these test can be optimized. Also command line option can be added to skip them refer https://pytest.org/en/stable/example/simple.html#control-skipping-of-tests-according-to-command-line-option

* Fixing test
2020-10-26 19:19:10 +01:00

48 lines
2.0 KiB
Python

from haystack import Finder
import pytest
@pytest.mark.slow
@pytest.mark.parametrize("retriever_with_docs", ["tfidf"], indirect=True)
def test_finder_get_answers(reader, retriever_with_docs, document_store_with_docs):
finder = Finder(reader, retriever_with_docs)
prediction = finder.get_answers(question="Who lives in Berlin?", top_k_retriever=10,
top_k_reader=3)
assert prediction is not None
assert prediction["question"] == "Who lives in Berlin?"
assert prediction["answers"][0]["answer"] == "Carla"
assert prediction["answers"][0]["probability"] <= 1
assert prediction["answers"][0]["probability"] >= 0
assert prediction["answers"][0]["meta"]["meta_field"] == "test1"
assert prediction["answers"][0]["context"] == "My name is Carla and I live in Berlin"
assert len(prediction["answers"]) == 3
@pytest.mark.parametrize("retriever_with_docs", ["tfidf"], indirect=True)
def test_finder_offsets(reader, retriever_with_docs, document_store_with_docs):
finder = Finder(reader, retriever_with_docs)
prediction = finder.get_answers(question="Who lives in Berlin?", top_k_retriever=10,
top_k_reader=5)
assert prediction["answers"][0]["offset_start"] == 11
assert prediction["answers"][0]["offset_end"] == 16
start = prediction["answers"][0]["offset_start"]
end = prediction["answers"][0]["offset_end"]
assert prediction["answers"][0]["context"][start:end] == prediction["answers"][0]["answer"]
@pytest.mark.slow
@pytest.mark.parametrize("retriever_with_docs", ["tfidf"], indirect=True)
def test_finder_get_answers_single_result(reader, retriever_with_docs, document_store_with_docs):
finder = Finder(reader, retriever_with_docs)
query = "testing finder"
prediction = finder.get_answers(question=query, top_k_retriever=1,
top_k_reader=1)
assert prediction is not None
assert len(prediction["answers"]) == 1