change test

This commit is contained in:
Julius Unverfehrt 2022-03-21 10:57:03 +01:00
parent 58ca784d6c
commit 3d9c4d8856

View File

@ -1,12 +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 = list(predictions) predictions = [p for p in predictions]
predictions = dict(**predictions[0]) assert predictions[0] == {'class': 'formula',
assert predictions == {'class': 'formula', 'probabilities': {'formula': 1.0, 'other': 1.576210689757148e-17,
'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