This commit is contained in:
Julius Unverfehrt 2022-03-21 11:01:21 +01:00
parent 3d9c4d8856
commit 3e9bfac5cf

View File

@ -1,11 +1,11 @@
def test_predict_pdf_works(predictor, test_pdf): def test_predict_pdf_works(predictor, test_pdf):
# FIXME strange outputs, with '\n's in the dict # FIXME strange outputs, with '\n's in the dict
predictions, metadata = predictor.predict_pdf(test_pdf) predictions, metadata = predictor.predict_pdf(test_pdf)
predictions = [p for p in predictions] predictions = [p for p in predictions][0]
assert predictions[0] == {'class': 'formula', assert predictions["class"] == "formula"
'probabilities': {'formula': 1.0, 'other': 1.576210689757148e-17, assert dict(**predictions["probabilities"]) == {'formula': 1.0, 'other': 1.576210689757148e-17,
'signature': 5.364939675629517e-24, 'signature': 5.364939675629517e-24,
'logo': 8.815339177836606e-25}} 'logo': 8.815339177836606e-25}
metadata = list(metadata) metadata = list(metadata)
metadata = dict(**metadata[0]) metadata = dict(**metadata[0])
metadata.pop("document_filename") # temp filename cannot be tested metadata.pop("document_filename") # temp filename cannot be tested