Leonid Pershin
aca8d5e990
- Updated the `_log_default` function to handle messages starting with `\r` for in-place updates. - Introduced `split_ssh_stream` to differentiate between line and progress events in SSH output. - Implemented `feed_ssh_log` to dispatch log messages based on event type. - Refactored `_stream_pty_output` to utilize the new logging and event handling functions. - Enhanced the `log` function in `term.py` to support progress updates, ensuring a cleaner output experience. - Updated `DownloadProgress` classes in `civitai_fetch.py` and `llamacpp_fetch.py` to print progress in place, improving user feedback during downloads.
34 lines
918 B
Python
34 lines
918 B
Python
"""SSH stream splitting for in-place progress bars."""
|
|
|
|
from gpu_rent.ssh_ops import feed_ssh_log, split_ssh_stream
|
|
|
|
|
|
def test_split_crlf_line():
|
|
events, rest = split_ssh_stream("hello\r\n")
|
|
assert events == [("line", "hello")]
|
|
assert rest == ""
|
|
|
|
|
|
def test_split_progress_cr():
|
|
events, rest = split_ssh_stream("prog 10%\rprog 20%\r")
|
|
assert events == [("progress", "prog 10%"), ("progress", "prog 20%")]
|
|
assert rest == ""
|
|
|
|
|
|
def test_split_mixed_progress_then_line():
|
|
events, rest = split_ssh_stream("a\rb\r\nok\npartial")
|
|
assert events == [
|
|
("progress", "a"),
|
|
("line", "b"),
|
|
("line", "ok"),
|
|
]
|
|
assert rest == "partial"
|
|
|
|
|
|
def test_feed_progress_prefix():
|
|
seen = []
|
|
feed_ssh_log(seen.append, "progress", "50%")
|
|
assert seen == ["\r50%"]
|
|
feed_ssh_log(seen.append, "line", "done")
|
|
assert seen[-1] == "done"
|