Function test_load_runtime

tests/test_runtime.py:117–162

(mocker: MockerFixture, tmp_path, runtime_type, expected_opts)

    ],
)
def test_load_runtime(mocker: MockerFixture, tmp_path, runtime_type, expected_opts):
    # mock model path
    model_path = pathlib.Path(tmp_path) / "fakeref" / "model.onnx"
    model_path.parent.mkdir(parents=True, exist_ok=True)
    model_path.touch()
    model_path = model_path.as_posix()

    # mock model metadata
    mock_model_metadata = MagicMock(spec=ModelInfo)

    # mock opts
    if runtime_type == RuntimeType.TORCHSCRIPT_32:
        mocker.patch("focoos.infer.runtimes.load_runtime.TORCH_AVAILABLE", True)
        mock_runtime_class = mocker.patch("focoos.infer.runtimes.torchscript.TorchscriptRuntime", autospec=True)
        mock_runtime_class.return_value = MagicMock(spec=TorchscriptRuntime, opts=expected_opts)
    else:
        mocker.patch("focoos.infer.runtimes.load_runtime.ORT_AVAILABLE", True)
        mock_runtime_class = mocker.patch("focoos.infer.runtimes.onnx.ONNXRuntime", autospec=True)
        mock_runtime_class.return_value = MagicMock(spec=ONNXRuntime, opts=expected_opts)

    # warmup_iter
    warmup_iter = 2

    # call the function to test
    runtime = load_runtime(
        runtime_type=runtime_type,
        model_path=model_path,
        model_info=mock_model_metadata,
        warmup_iter=warmup_iter,
    )

    # assertions
    assert runtime is not None
    if runtime_type == RuntimeType.TORCHSCRIPT_32:
        mock_runtime_class.assert_called_once_with(
            model_path=model_path,
            opts=expected_opts,
            model_info=mock_model_metadata,
            device="auto",
        )
    else:
        mock_runtime_class.assert_called_once_with(
            model_path,
            expected_opts,
            mock_model_metadata,
        )


def test_load_unavailable_runtime(mocker: MockerFixture):