add predictor test

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()

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}