From 02829ef986bc2a5c4f33e9c45c9267bcf2d07a1d Mon Sep 17 00:00:00 2001 From: SWHL Date: Tue, 21 May 2024 20:52:17 +0800 Subject: [PATCH] test(rapidocr_onnxruntime): Add two unit testings about cuda and dml --- python/tests/test_ort.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) 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)