chore: import upstream snapshot with attribution

python/ray/data/tests/test_progress_bar.py

@@ -0,0 +1,134 @@
+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__]))