diff --git a/test/conftest.py b/test/conftest.py index ce65f3d..73f2e37 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,5 +1,9 @@ +import os.path + import pytest +from image_prediction.predictor import Predictor + @pytest.fixture def predictions(): @@ -53,3 +57,14 @@ def response(): "position": {"pageNumber": 9, "x1": 63.5, "x2": 202.85000000000002, "y1": 472.0, "y2": 533.05}, } ] + + +@pytest.fixture +def predictor(): + return Predictor() + + +@pytest.fixture +def test_pdf(): + with open("./test/test_data/f2dc689ca794fccb8cd38b95f2bf6ba9.pdf", "rb") as f: + return f.read() \ No newline at end of file diff --git a/test/test_data/f2dc689ca794fccb8cd38b95f2bf6ba9.pdf b/test/test_data/f2dc689ca794fccb8cd38b95f2bf6ba9.pdf new file mode 100644 index 0000000..41f0d70 Binary files /dev/null and b/test/test_data/f2dc689ca794fccb8cd38b95f2bf6ba9.pdf differ diff --git a/test/unit_tests/test_predictor.py b/test/unit_tests/test_predictor.py new file mode 100644 index 0000000..0523adf --- /dev/null +++ b/test/unit_tests/test_predictor.py @@ -0,0 +1,11 @@ +def test_predict_pdf_works(predictor, test_pdf): + predictions, metadata = predictor.predict_pdf(test_pdf) + assert list(predictions) == [{'class': 'formula', + 'probabilities': {'formula': 1.0, 'other': 1.576210689757148e-17, + 'signature': 5.364939675629517e-24, + 'logo': 8.815339177836606e-25}}] + # FIXME strange output for metadata, with '\n's in the dict + metadata = list(metadata) + metadata = dict(**metadata[0]) + metadata.pop("document_filename") # temp filename cannot be tested + assert metadata == {'px_width': 389.0, 'px_height': 389.0, 'width': 194.49999000000003, 'height': 194.49998999999997, 'x1': 320.861, 'x2': 515.36099, 'y1': 347.699, 'y2': 542.19899, 'page_width': 595.2800000000001, 'page_height': 841.89, 'page_rotation': 0, 'page_idx': 1, 'n_pages': 3}