Fix tests progressbar

brainets · Oct 29, 2024 · feb9eb7 · feb9eb7
1 parent 38cc6bc
commit feb9eb7
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/hoi/utils/tests/test_progressbar.py b/hoi/utils/tests/test_progressbar.py
@@ -1,6 +1,8 @@
-from jax_tqdm import scan_tqdm, loop_tqdm, build_tqdm
+from jax_tqdm import scan_tqdm, loop_tqdm
+from hoi.utils.progressbar import get_pbar
 import typing
 import pytest
+from tqdm.auto import tqdm
 
 
 class TestProgressBar(object):
@@ -12,8 +14,6 @@ def test_scan_tqdm(self, print_rate):
     def loop_scan_tqdm(self, print_rate):
         assert isinstance(loop_tqdm(10, print_rate), typing.Callable)
 
-    @pytest.mark.parametrize("print_rate", [2, 10])
-    def test_build_tqdm(self, print_rate):
-        val = build_tqdm(10, print_rate)
-        assert isinstance(val, tuple)
-        assert len(val) == 2
+    def test_get_pbar(self):
+        """Test function get_pbar."""
+        assert isinstance(get_pbar(iterable=range(10), leave=False), tqdm)