wehub-resource-sync
135 lines
3.8 KiB
Python
135 lines
3.8 KiB
Python
import functools
|
|
import logging
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from pytest import fixture
|
|
|
|
import ray
|
|
from ray.data._internal.progress.progress_bar import ProgressBar
|
|
|
|
|
|
@fixture(params=[True, False])
|
|
def enable_tqdm_ray(request):
|
|
context = ray.data.DataContext.get_current()
|
|
original_use_ray_tqdm = context.use_ray_tqdm
|
|
context.use_ray_tqdm = request.param
|
|
yield request.param
|
|
context.use_ray_tqdm = original_use_ray_tqdm
|
|
|
|
|
|
def test_progress_bar(enable_tqdm_ray):
|
|
total = 10
|
|
# Used to record the total value of the bar at close
|
|
total_at_close = 0
|
|
|
|
def patch_close(bar):
|
|
nonlocal total_at_close
|
|
total_at_close = 0
|
|
original_close = bar.close
|
|
|
|
@functools.wraps(original_close)
|
|
def wrapped_close():
|
|
nonlocal total_at_close
|
|
total_at_close = bar.total
|
|
|
|
bar.close = wrapped_close
|
|
|
|
# Test basic usage
|
|
pb = ProgressBar("", total, "unit", enabled=True)
|
|
assert pb._bar is not None
|
|
patch_close(pb._bar)
|
|
for _ in range(total):
|
|
pb.update(1)
|
|
pb.close()
|
|
|
|
assert pb._progress == total
|
|
assert total_at_close == total
|
|
|
|
# Test if update() exceeds the original total, the total will be updated.
|
|
pb = ProgressBar("", total, "unit", enabled=True)
|
|
assert pb._bar is not None
|
|
patch_close(pb._bar)
|
|
new_total = total * 2
|
|
for _ in range(new_total):
|
|
pb.update(1)
|
|
pb.close()
|
|
|
|
assert pb._progress == new_total
|
|
assert total_at_close == new_total
|
|
|
|
# Test that if the bar is not complete at close(), the total will be updated.
|
|
pb = ProgressBar("", total, "unit")
|
|
assert pb._bar is not None
|
|
patch_close(pb._bar)
|
|
new_total = total // 2
|
|
for _ in range(new_total):
|
|
pb.update(1)
|
|
pb.close()
|
|
|
|
assert pb._progress == new_total
|
|
assert total_at_close == new_total
|
|
|
|
# Test updating the total
|
|
pb = ProgressBar("", total, "unit", enabled=True)
|
|
assert pb._bar is not None
|
|
patch_close(pb._bar)
|
|
new_total = total * 2
|
|
pb.update(0, new_total)
|
|
|
|
assert pb._bar.total == new_total
|
|
pb.update(total + 1, total)
|
|
assert pb._bar.total == total + 1
|
|
pb.close()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"name, expected_description, max_line_length, should_emit_warning",
|
|
[
|
|
("Op", "Op", 2, False),
|
|
("Op->Op", "Op->Op", 5, False),
|
|
("Op->Op->Op", "Op->...->Op", 9, True),
|
|
("Op->Op->Op", "Op->Op->Op", 10, False),
|
|
# Test case for https://github.com/ray-project/ray/issues/47679.
|
|
("spam", "spam", 1, False),
|
|
],
|
|
)
|
|
def test_progress_bar_truncates_chained_operators(
|
|
name,
|
|
expected_description,
|
|
max_line_length,
|
|
should_emit_warning,
|
|
caplog,
|
|
propagate_logs,
|
|
):
|
|
with patch.object(ProgressBar, "MAX_NAME_LENGTH", max_line_length):
|
|
pb = ProgressBar(name, None, "unit")
|
|
|
|
assert pb.get_description() == expected_description
|
|
if should_emit_warning:
|
|
assert any(
|
|
record.levelno == logging.WARNING
|
|
and "Truncating long operator name" in record.message
|
|
for record in caplog.records
|
|
), caplog.records
|
|
|
|
|
|
def test_progress_bar_tqdm_not_installed(enable_tqdm_ray):
|
|
"""Test behavior when tqdm package is not installed."""
|
|
with patch("sys.stdout.isatty", return_value=True):
|
|
# Mock tqdm not being installed
|
|
with patch("ray.data._internal.progress.progress_bar.tqdm", None):
|
|
pb = ProgressBar("test", 100, "unit", enabled=True)
|
|
if enable_tqdm_ray:
|
|
# tqdm_ray still works (part of Ray, no tqdm dependency)
|
|
assert pb._bar is not None
|
|
else:
|
|
# No progress bar available (interactive, so no logging fallback)
|
|
assert pb._bar is None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|