aboutsummaryrefslogtreecommitdiff
path: root/tests/test_utils_logging.py
diff options
context:
space:
mode:
authorDmitrii Agibov <dmitrii.agibov@arm.com>2022-10-24 15:08:08 +0100
committerDmitrii Agibov <dmitrii.agibov@arm.com>2022-10-26 17:08:13 +0100
commit58a65fee574c00329cf92b387a6d2513dcbf6100 (patch)
tree47e3185f78b4298ab029785ddee68456e44cac10 /tests/test_utils_logging.py
parent9d34cb72d45a6d0a2ec1063ebf32536c1efdba75 (diff)
downloadmlia-58a65fee574c00329cf92b387a6d2513dcbf6100.tar.gz
MLIA-433 Add TensorFlow Lite compatibility check
- Add ability to intercept low level TensorFlow output - Produce advice for the models that could not be converted to the TensorFlow Lite format - Refactor utility functions for TensorFlow Lite conversion - Add TensorFlow Lite compatibility checker Change-Id: I47d120d2619ced7b143bc92c5184515b81c0220d
Diffstat (limited to 'tests/test_utils_logging.py')
-rw-r--r--tests/test_utils_logging.py24
1 files changed, 23 insertions, 1 deletions
diff --git a/tests/test_utils_logging.py b/tests/test_utils_logging.py
index 1e212b2..ac835c6 100644
--- a/tests/test_utils_logging.py
+++ b/tests/test_utils_logging.py
@@ -8,10 +8,14 @@ import sys
from contextlib import ExitStack as does_not_raise
from pathlib import Path
from typing import Any
+from typing import Callable
+from unittest.mock import MagicMock
import pytest
-from mlia.cli.logging import create_log_handler
+from mlia.utils.logging import create_log_handler
+from mlia.utils.logging import redirect_output
+from mlia.utils.logging import redirect_raw_output
@pytest.mark.parametrize(
@@ -62,3 +66,21 @@ def test_create_log_handler(
delay=delay,
)
assert isinstance(handler, expected_class)
+
+
+@pytest.mark.parametrize(
+ "redirect_context_manager",
+ [
+ redirect_raw_output,
+ redirect_output,
+ ],
+)
+def test_output_redirection(redirect_context_manager: Callable) -> None:
+ """Test output redirection via context manager."""
+ print("before redirect")
+ logger_mock = MagicMock()
+ with redirect_context_manager(logger_mock):
+ print("output redirected")
+ print("after redirect")
+
+ logger_mock.log.assert_called_once_with(logging.INFO, "output redirected")