update tests

This commit is contained in:
mertalev 2025-03-14 16:07:35 -04:00
parent fedbecc9f2
commit 1785d0916b
No known key found for this signature in database
GPG Key ID: F7C271C07CF04AAE
2 changed files with 13 additions and 17 deletions

View File

@ -22,7 +22,7 @@ def get_soc(device_tree_path: Path | str) -> str | None:
return soc
log.warning("Device is not supported for RKNN")
except OSError as e:
log.warning("Could not read /proc/device-tree/compatible. Reason: %s", e)
log.warning(f"Could not read {device_tree_path}. Reason: %s", e)
return None
@ -43,7 +43,7 @@ def init_rknn(model_path: str) -> "RKNNLite":
rknn_lite = RKNNLite()
ret = rknn_lite.load_rknn(model_path)
if ret != 0:
raise RuntimeError("Load RKNN rknnModel failed")
raise RuntimeError("Failed to load RKNN model")
if soc_name in RKNN_COREMASK_SUPPORTED_SOCS:
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_AUTO)
@ -51,7 +51,7 @@ def init_rknn(model_path: str) -> "RKNNLite":
ret = rknn_lite.init_runtime() # Please do not set this parameter on other platforms.
if ret != 0:
raise RuntimeError("Init runtime environment failed")
raise RuntimeError("Failed to inititalize RKNN runtime environment")
return rknn_lite

View File

@ -70,13 +70,13 @@ class TestBase:
assert encoder.model_format == ModelFormat.ARMNN
def test_sets_default_model_format_rknn(self, mocker: MockerFixture) -> None:
def test_sets_default_model_format_to_rknn_if_available(self, mocker: MockerFixture) -> None:
mocker.patch.object(settings, "rknn", True)
mocker.patch("rknn.rknnpool.is_available", False)
mocker.patch("app.sessions.rknn.is_available", True)
encoder = OpenClipTextualEncoder("ViT-B-32__openai")
assert encoder.model_format == ModelFormat.ONNX
assert encoder.model_format == ModelFormat.RKNN
def test_casts_cache_dir_string_to_path(self) -> None:
cache_dir = "/test_cache"
@ -353,15 +353,12 @@ class TestRknnSession:
model_path = mock.MagicMock(spec=Path)
tpe = 1
mocker.patch("app.sessions.rknn.soc_name", "rk3566")
mocker.patch("app.sessions.rknn.is_available", True)
RknnSession(model_path)
rknn_session.assert_called_once_with(
rknnModel=Path(str(model_path).replace("model", "rk3566")).as_posix(), tpes=tpe, func=run_inference
)
rknn_session.assert_called_once_with(model_path=model_path.as_posix(), tpes=tpe, func=run_inference)
info.assert_has_calls(
[mock.call(f"Loaded RKNN model from {str(model_path).replace('model', 'rk3566')} with {tpe} threads.")]
)
info.assert_has_calls([mock.call(f"Loaded RKNN model from {model_path} with {tpe} threads.")])
def test_run_rknn(self, rknn_session: mock.Mock, mocker: MockerFixture) -> None:
rknn_session.return_value.load.return_value = 123
@ -879,9 +876,7 @@ class TestLoad:
mock_model.clear_cache.assert_not_called()
mock_model.load.assert_not_called()
async def test_falls_back_to_onnx_if_other_format_does_not_exist(
self, exception: mock.Mock, warning: mock.Mock
) -> None:
async def test_falls_back_to_onnx_if_other_format_does_not_exist(self, warning: mock.Mock) -> None:
mock_model = mock.Mock(spec=InferenceModel)
mock_model.model_name = "test_model_name"
mock_model.model_type = ModelType.VISUAL
@ -896,8 +891,9 @@ class TestLoad:
mock_model.clear_cache.assert_not_called()
assert mock_model.load.call_count == 2
exception.assert_called_once_with(error)
warning.assert_called_once_with("ARMNN is available, but model 'test_model_name' does not support it.")
warning.assert_called_once_with(
"ARMNN is available, but model 'test_model_name' does not support it.", exc_info=error
)
mock_model.model_format = ModelFormat.ONNX