mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-10-31 01:39:45 +00:00
165a0a5faa
* mock all translator tests and move one to e2e * typo * extract pipeline tests using translator * remove duplicate test * move generator test in e2e * Update e2e/pipelines/test_extractive_qa.py * pytest.mark.unit * black * remove model name as well * remove unused fixture * rename original and improve pipeline tests * fixes * pylint
116 lines
3.3 KiB
Python
116 lines
3.3 KiB
Python
import pytest
|
|
|
|
import haystack
|
|
from haystack.schema import Document
|
|
from haystack.nodes import TransformersTranslator
|
|
|
|
|
|
ORIGINAL_TEXT = "TEST QUERY"
|
|
TRANSLATION = "MOCK TRANSLATION"
|
|
|
|
|
|
class MockTokenizer:
|
|
@classmethod
|
|
def from_pretrained(cls, *a, **k):
|
|
return cls()
|
|
|
|
def __call__(self, *a, **k):
|
|
return self
|
|
|
|
def to(self, *a, **k):
|
|
return {}
|
|
|
|
def batch_decode(self, *a, **k):
|
|
return [TRANSLATION]
|
|
|
|
|
|
class MockModel:
|
|
@classmethod
|
|
def from_pretrained(cls, *a, **k):
|
|
return cls()
|
|
|
|
def generate(self, *a, **k):
|
|
return None
|
|
|
|
def to(self, *a, **k):
|
|
return None
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_models(monkeypatch):
|
|
monkeypatch.setattr(haystack.nodes.translator.transformers, "AutoModelForSeq2SeqLM", MockModel)
|
|
monkeypatch.setattr(haystack.nodes.translator.transformers, "AutoTokenizer", MockTokenizer)
|
|
|
|
|
|
@pytest.fixture
|
|
def en_to_de_translator(mock_models) -> TransformersTranslator:
|
|
return TransformersTranslator(model_name_or_path="irrelevant/anyway")
|
|
|
|
|
|
@pytest.fixture
|
|
def de_to_en_translator(mock_models) -> TransformersTranslator:
|
|
return TransformersTranslator(model_name_or_path="irrelevant/anyway")
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_translator_with_query(en_to_de_translator):
|
|
assert en_to_de_translator.translate(query=ORIGINAL_TEXT) == TRANSLATION
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_translator_with_list(en_to_de_translator):
|
|
assert en_to_de_translator.translate(documents=[ORIGINAL_TEXT])[0] == TRANSLATION
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_translator_with_document(en_to_de_translator):
|
|
assert en_to_de_translator.translate(documents=[Document(content=ORIGINAL_TEXT)])[0].content == TRANSLATION
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_translator_with_document_preserves_original(en_to_de_translator):
|
|
original_document = Document(content=ORIGINAL_TEXT)
|
|
en_to_de_translator.translate(documents=[original_document])
|
|
assert original_document.content == ORIGINAL_TEXT
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_translator_with_dictionary(en_to_de_translator):
|
|
assert en_to_de_translator.translate(documents=[{"content": ORIGINAL_TEXT}])[0]["content"] == TRANSLATION
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_translator_with_dictionary_preserves_original(en_to_de_translator):
|
|
original_document = {"content": ORIGINAL_TEXT}
|
|
en_to_de_translator.translate(documents=[original_document])
|
|
assert original_document["content"] == ORIGINAL_TEXT
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_translator_with_dictionary_with_dict_key(en_to_de_translator):
|
|
assert en_to_de_translator.translate(documents=[{"key": ORIGINAL_TEXT}], dict_key="key")[0]["key"] == TRANSLATION
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_translator_with_empty_original(en_to_de_translator):
|
|
with pytest.raises(AttributeError):
|
|
en_to_de_translator.translate()
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_translator_with_query_and_documents(en_to_de_translator):
|
|
with pytest.raises(AttributeError):
|
|
en_to_de_translator.translate(query=ORIGINAL_TEXT, documents=[ORIGINAL_TEXT])
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_translator_with_dict_without_text_key(en_to_de_translator):
|
|
with pytest.raises(AttributeError):
|
|
en_to_de_translator.translate(documents=[{"text1": ORIGINAL_TEXT}])
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_translator_with_dict_with_non_string_value(en_to_de_translator):
|
|
with pytest.raises(AttributeError):
|
|
en_to_de_translator.translate(documents=[{"text": 123}])
|