rapidocr/python/tests/test_engine_config.py

# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
import sys
from pathlib import Path

import pytest

root_dir = Path(__file__).resolve().parent.parent
sys.path.append(str(root_dir))

from rapidocr.inference_engine.onnxruntime import provider_config
from rapidocr.inference_engine.onnxruntime.provider_config import EP, ProviderConfig
from rapidocr.inference_engine.openvino.device_config import CPUConfig


class AttrDict(dict):
    def __getattr__(self, key):
        try:
            return self[key]
        except KeyError as e:
            raise AttributeError(key) from e

    def __setattr__(self, key, value):
        self[key] = value


def make_provider_cfg(**overrides):
    cfg = AttrDict(
        use_cuda=False,
        cuda_ep_cfg={"device_id": 0},
        use_dml=False,
        dml_ep_cfg=None,
        use_cann=False,
        cann_ep_cfg={"device_id": 0},
        use_coreml=False,
        coreml_ep_cfg={"ModelFormat": "MLProgram"},
        cpu_ep_cfg={"arena_extend_strategy": "kSameAsRequested"},
    )
    cfg.update(overrides)
    return cfg


def make_provider_config(monkeypatch, cfg, providers, device="CPU"):
    monkeypatch.setattr(provider_config, "get_available_providers", lambda: providers)
    monkeypatch.setattr(provider_config, "get_device", lambda: device)
    return ProviderConfig(cfg)


def test_onnxruntime_provider_config_uses_cpu_by_default(monkeypatch):
    cfg = make_provider_cfg()
    provider_cfg = make_provider_config(
        monkeypatch,
        cfg,
        providers=[EP.CPU_EP.value],
    )

    assert provider_cfg.get_ep_list() == [(EP.CPU_EP.value, cfg.cpu_ep_cfg)]


def test_onnxruntime_provider_config_prefers_cuda_when_available(monkeypatch):
    cfg = make_provider_cfg(use_cuda=True)
    provider_cfg = make_provider_config(
        monkeypatch,
        cfg,
        providers=[EP.CUDA_EP.value, EP.CPU_EP.value],
        device="GPU",
    )

    assert provider_cfg.get_ep_list() == [
        (EP.CUDA_EP.value, cfg.cuda_ep_cfg),
        (EP.CPU_EP.value, cfg.cpu_ep_cfg),
    ]


def test_onnxruntime_provider_config_falls_back_when_cuda_unavailable(monkeypatch):
    cfg = make_provider_cfg(use_cuda=True)
    provider_cfg = make_provider_config(
        monkeypatch,
        cfg,
        providers=[EP.CPU_EP.value],
        device="CPU",
    )

    assert provider_cfg.get_ep_list() == [(EP.CPU_EP.value, cfg.cpu_ep_cfg)]


def test_onnxruntime_provider_config_prefers_directml_on_windows(monkeypatch):
    cfg = make_provider_cfg(use_dml=True, dml_ep_cfg={"device_id": 0})
    provider_cfg = make_provider_config(
        monkeypatch,
        cfg,
        providers=[EP.DIRECTML_EP.value, EP.CPU_EP.value],
    )
    monkeypatch.setattr(provider_config.platform, "system", lambda: "Windows")
    monkeypatch.setattr(provider_config.platform, "version", lambda: "10.0.19045")

    assert provider_cfg.get_ep_list() == [
        (EP.DIRECTML_EP.value, cfg.dml_ep_cfg),
        (EP.CPU_EP.value, cfg.cpu_ep_cfg),
    ]


def test_onnxruntime_provider_config_skips_directml_on_non_windows(monkeypatch):
    cfg = make_provider_cfg(use_dml=True)
    provider_cfg = make_provider_config(
        monkeypatch,
        cfg,
        providers=[EP.DIRECTML_EP.value, EP.CPU_EP.value],
    )
    monkeypatch.setattr(provider_config.platform, "system", lambda: "Linux")

    assert provider_cfg.get_ep_list() == [(EP.CPU_EP.value, cfg.cpu_ep_cfg)]


def test_onnxruntime_provider_config_adds_cann_when_available(monkeypatch):
    cfg = make_provider_cfg(use_cann=True)
    provider_cfg = make_provider_config(
        monkeypatch,
        cfg,
        providers=[EP.CANN_EP.value, EP.CPU_EP.value],
    )

    assert provider_cfg.get_ep_list() == [
        (EP.CANN_EP.value, cfg.cann_ep_cfg),
        (EP.CPU_EP.value, cfg.cpu_ep_cfg),
    ]


def test_onnxruntime_provider_config_adds_coreml_on_darwin(monkeypatch):
    cfg = make_provider_cfg(use_coreml=True)
    provider_cfg = make_provider_config(
        monkeypatch,
        cfg,
        providers=[EP.COREML_EP.value, EP.CPU_EP.value],
    )
    monkeypatch.setattr(provider_config.platform, "system", lambda: "Darwin")

    assert provider_cfg.get_ep_list() == [
        (EP.COREML_EP.value, cfg.coreml_ep_cfg),
        (EP.CPU_EP.value, cfg.cpu_ep_cfg),
    ]


def test_onnxruntime_provider_config_rejects_empty_session_providers(monkeypatch):
    cfg = make_provider_cfg()
    provider_cfg = make_provider_config(
        monkeypatch,
        cfg,
        providers=[EP.CPU_EP.value],
    )

    with pytest.raises(ValueError, match="Session Providers is empty"):
        provider_cfg.verify_providers([])


def test_openvino_cpu_config_omits_defaults():
    cfg = AttrDict(
        inference_num_threads=-1,
        performance_hint=None,
        performance_num_requests=-1,
        enable_cpu_pinning=None,
        num_streams=-1,
        enable_hyper_threading=None,
        scheduling_core_type=None,
    )

    assert CPUConfig(cfg).get_config() == {}


def test_openvino_cpu_config_maps_non_default_values(monkeypatch):
    cfg = AttrDict(
        inference_num_threads=4,
        performance_hint="LATENCY",
        performance_num_requests=2,
        enable_cpu_pinning=True,
        num_streams=1,
        enable_hyper_threading=False,
        scheduling_core_type="PCORE_ONLY",
    )
    monkeypatch.setattr("rapidocr.inference_engine.openvino.device_config.os.cpu_count", lambda: 8)

    assert CPUConfig(cfg).get_config() == {
        "INFERENCE_NUM_THREADS": "4",
        "PERFORMANCE_HINT": "LATENCY",
        "PERFORMANCE_HINT_NUM_REQUESTS": "2",
        "ENABLE_CPU_PINNING": "True",
        "NUM_STREAMS": "1",
        "ENABLE_HYPER_THREADING": "False",
        "SCHEDULING_CORE_TYPE": "PCORE_ONLY",
    }


def test_openvino_cpu_config_ignores_invalid_thread_count(monkeypatch):
    cfg = AttrDict(
        inference_num_threads=16,
        performance_hint=None,
        performance_num_requests=-1,
        enable_cpu_pinning=None,
        num_streams=-1,
        enable_hyper_threading=None,
        scheduling_core_type=None,
    )
    monkeypatch.setattr("rapidocr.inference_engine.openvino.device_config.os.cpu_count", lambda: 8)

    assert "INFERENCE_NUM_THREADS" not in CPUConfig(cfg).get_config()