fix up test_utils

tier4 · Nov 28, 2024 · aaa1984 · aaa1984
1 parent f1be14b
commit aaa1984
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/src/otaclient/_utils.py b/src/otaclient/_utils.py
@@ -23,7 +23,7 @@
 import time
 import traceback
 from pathlib import Path
-from typing import Callable
+from typing import Callable, Literal
 
 from otaclient._types import OTAClientStatus
 from otaclient_common._io import read_str_from_file, write_str_to_file_atomic
@@ -37,7 +37,7 @@ def wait_and_log(
     check_flag: Callable[[], bool],
     message: str = "",
     *,
-    check_for: bool = True,
+    check_for: Literal[True] | Literal[False] = True,
     check_interval: int = 2,
     log_interval: int = 30,
     log_func: Callable[[str], None] = logger.info,

diff --git a/tests/test_otaclient/test_utils.py b/tests/test_otaclient/test_utils.py
@@ -42,8 +42,9 @@ def test_wait_and_log(caplog: pytest.LogCaptureFixture):
     _msg = "ticking flag"
 
     wait_and_log(
-        _flag,
+        _flag.is_set,
         _msg,
+        check_for=True,
         check_interval=1,
         log_interval=2,
         log_func=logger.warning,