haystack/test/test_file_converter.py
bogdankostic 834f8c4902
Change return types of indexing pipeline nodes (#2342)
* Change return types of file converters

* Change return types of preprocessor

* Change return types of crawler

* Adapt utils to functions to new return types

* Adapt __init__.py to new method names

* Prevent circular imports

* Update Documentation & Code Style

* Let DocStores' run method accept Documents

* Adapt tests to new return types

* Update Documentation & Code Style

* Put "# type: ignore" to right place

* Remove id_hash_keys property from Document primitive

* Update Documentation & Code Style

* Adapt tests to new return types and missing id_hash_keys property

* Fix mypy

* Fix mypy

* Adapt PDFToTextOCRConverter

* Remove id_hash_keys from RestAPI tests

* Update Documentation & Code Style

* Rename tests

* Remove redundant setting of content_type="text"

* Add DeprecationWarning

* Add id_hash_keys to elasticsearch_index_to_document_store

* Change document type from dict to Docuemnt in PreProcessor test

* Fix file path in Tutorial 5

* Remove added output in Tutorial 5

* Update Documentation & Code Style

* Fix file_paths in Tutorial 9 + fix gz files in fetch_archive_from_http

* Adapt tutorials to new return types

* Adapt tutorial 14 to new return types

* Update Documentation & Code Style

* Change assertions to HaystackErrors

* Import HaystackError correctly

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
2022-03-29 13:53:35 +02:00

122 lines
4.9 KiB
Python

from pathlib import Path
import os
import pytest
from haystack.nodes import (
MarkdownConverter,
DocxToTextConverter,
PDFToTextConverter,
PDFToTextOCRConverter,
TikaConverter,
AzureConverter,
ParsrConverter,
)
from .conftest import SAMPLES_PATH
@pytest.mark.tika
@pytest.mark.parametrize(
# "Converter", [PDFToTextConverter, TikaConverter, PDFToTextOCRConverter]
"Converter",
[PDFToTextOCRConverter],
)
def test_convert(Converter):
converter = Converter()
document = converter.convert(file_path=SAMPLES_PATH / "pdf" / "sample_pdf_1.pdf")[0]
pages = document.content.split("\f")
assert len(pages) == 4 # the sample PDF file has four pages.
assert pages[0] != "" # the page 1 of PDF contains text.
assert pages[2] == "" # the page 3 of PDF file is empty.
# assert text is retained from the document.
# As whitespace can differ (\n," ", etc.), we standardize all to simple whitespace
page_standard_whitespace = " ".join(pages[0].split())
assert "Adobe Systems made the PDF specification available free of charge in 1993." in page_standard_whitespace
@pytest.mark.tika
@pytest.mark.parametrize("Converter", [PDFToTextConverter, TikaConverter])
def test_table_removal(Converter):
converter = Converter(remove_numeric_tables=True)
document = converter.convert(file_path=SAMPLES_PATH / "pdf" / "sample_pdf_1.pdf")[0]
pages = document.content.split("\f")
# assert numeric rows are removed from the table.
assert "324" not in pages[0]
assert "54x growth" not in pages[0]
@pytest.mark.tika
@pytest.mark.parametrize("Converter", [PDFToTextConverter, TikaConverter])
def test_language_validation(Converter, caplog):
converter = Converter(valid_languages=["en"])
converter.convert(file_path=SAMPLES_PATH / "pdf" / "sample_pdf_1.pdf")
assert "samples/pdf/sample_pdf_1.pdf is not one of ['en']." not in caplog.text
converter = Converter(valid_languages=["de"])
converter.convert(file_path=SAMPLES_PATH / "pdf" / "sample_pdf_1.pdf")
assert "samples/pdf/sample_pdf_1.pdf is not one of ['de']." in caplog.text
def test_docx_converter():
converter = DocxToTextConverter()
document = converter.convert(file_path=SAMPLES_PATH / "docx" / "sample_docx.docx")[0]
assert document.content.startswith("Sample Docx File")
def test_markdown_converter():
converter = MarkdownConverter()
document = converter.convert(file_path=SAMPLES_PATH / "markdown" / "sample.md")[0]
assert document.content.startswith("What to build with Haystack")
def test_azure_converter():
# Check if Form Recognizer endpoint and credential key in environment variables
if "AZURE_FORMRECOGNIZER_ENDPOINT" in os.environ and "AZURE_FORMRECOGNIZER_KEY" in os.environ:
converter = AzureConverter(
endpoint=os.environ["AZURE_FORMRECOGNIZER_ENDPOINT"],
credential_key=os.environ["AZURE_FORMRECOGNIZER_KEY"],
save_json=True,
)
docs = converter.convert(file_path=SAMPLES_PATH / "pdf" / "sample_pdf_1.pdf")
assert len(docs) == 2
assert docs[0].content_type == "table"
assert docs[0].content.shape[0] == 4 # number of rows
assert docs[0].content.shape[1] == 5 # number of columns, Form Recognizer assumes there are 5 columns
assert list(docs[0].content.columns) == ["", "Column 1", "", "Column 2", "Column 3"]
assert list(docs[0].content.iloc[3]) == ["D", "$54.35", "", "$6345.", ""]
assert (
docs[0].meta["preceding_context"] == "specification. These proprietary technologies are not "
"standardized and their\nspecification is published only on "
"Adobe's website. Many of them are also not\nsupported by "
"popular third-party implementations of PDF."
)
assert docs[0].meta["following_context"] == ""
assert docs[1].content_type == "text"
assert docs[1].content.startswith("A sample PDF file")
def test_parsr_converter():
converter = ParsrConverter()
docs = converter.convert(file_path=str((SAMPLES_PATH / "pdf" / "sample_pdf_1.pdf").absolute()))
assert len(docs) == 2
assert docs[0].content_type == "table"
assert docs[0].content.shape[0] == 4 # number of rows
assert docs[0].content.shape[1] == 4
assert list(docs[0].content.columns) == ["", "Column 1", "Column 2", "Column 3"]
assert list(docs[0].content.iloc[3]) == ["D", "$54.35", "$6345.", ""]
assert (
docs[0].meta["preceding_context"] == "specification. These proprietary technologies are not "
"standardized and their\nspecification is published only on "
"Adobe's website. Many of them are also not\nsupported by popular "
"third-party implementations of PDF."
)
assert docs[0].meta["following_context"] == ""
assert docs[1].content_type == "text"
assert docs[1].content.startswith("A sample PDF file")
assert docs[1].content.endswith("Page 4 of Sample PDF\n… the page 3 is empty.")