2022-03-21 13:54:27 +01:00

27 lines
1.0 KiB
Python

def test_predict_pdf_works(predictor, test_pdf):
# FIXME ugly test since there are '\n's in the dict with unknown heritage
predictions, metadata = predictor.predict_pdf(test_pdf)
predictions = [p for p in predictions][0]
assert predictions["class"] == "formula"
probabilities = predictions["probabilities"]
# Floating point precision problem for output so test only that keys exist not the values
assert all(key in probabilities for key in ("formula", "other", "signature", "logo"))
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,
}