diff --git a/python/tests/test_ort.py b/python/tests/test_ort.py index 03f8e0206..b7304d4cc 100644 --- a/python/tests/test_ort.py +++ b/python/tests/test_ort.py @@ -1,6 +1,7 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com +import logging import sys from pathlib import Path @@ -20,6 +21,22 @@ package_name = "rapidocr_onnxruntime" +def test_ort_cuda_warning(caplog): + engine = RapidOCR(det_use_cuda=True) + caplog.set_level(logging.WARNING) + + assert caplog.records[0].levelname == "WARNING" + assert "CUDAExecutionProvider" in caplog.records[0].message + + +def test_ort_dml_warning(caplog): + engine = RapidOCR(det_use_dml=True) + caplog.set_level(logging.WARNING) + + assert caplog.records[0].levelname == "WARNING" + assert "DirectML" in caplog.records[0].message + + def test_mode_one_img(): img_path = tests_dir / "issue_170.png" result, _ = engine(img_path)