From 2ffc9c45b63873f42ebcf5d89fb9b7ddffa4b585 Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <y.t.d.huang@phys.ncku.edu.tw>
Date: Fri, 5 Apr 2024 17:11:24 +0800
Subject: [PATCH 1/2] Support general `I/O` for `ProgressBar` output

---
 src/progress_bar.jl       | 12 ++++++------
 test/low_rank_dynamics.jl |  4 +++-
 test/progress_bar.jl      | 14 ++++++++++++++
 test/runtests.jl          |  6 +++---
 4 files changed, 26 insertions(+), 10 deletions(-)
 create mode 100644 test/progress_bar.jl

diff --git a/src/progress_bar.jl b/src/progress_bar.jl
index b9e272fc..ed36755a 100644
--- a/src/progress_bar.jl
+++ b/src/progress_bar.jl
@@ -1,4 +1,4 @@
-export ProgressBar
+export ProgressBar, next!
 
 struct ProgressBar{CT,T1<:Integer, T2<:Real}
     counter::CT
@@ -9,11 +9,11 @@ struct ProgressBar{CT,T1<:Integer, T2<:Real}
     lock::ReentrantLock
 end
 
-function ProgressBar(max_counts; enable=true, bar_width=30)
+function ProgressBar(max_counts::Int; enable::Bool=true, bar_width::Int=30)
     return ProgressBar(Ref{Int64}(0), max_counts, enable, bar_width, time(), ReentrantLock())
 end
 
-function next!(p::ProgressBar)
+function next!(p::ProgressBar, io::IO=stdout)
 
     lock(p.lock)
 
@@ -43,10 +43,10 @@ function next!(p::ProgressBar)
     # Construct the progress bar string
     bar = "[" * repeat("=", progress) * repeat(" ", bar_width - progress) * "]"
 
-    print("\rProgress: $bar $percentage_100% --- Elapsed Time: $elapsed_time_str (ETA: $eta_str)")
-    flush(stdout)
+    print(io, "\rProgress: $bar $percentage_100% --- Elapsed Time: $elapsed_time_str (ETA: $eta_str)")
+    flush(io)
 
     unlock(p.lock)
 
-    p.counter[] >= p.max_counts ? print("\n") : nothing
+    p.counter[] >= p.max_counts ? print(io, "\n") : nothing
 end
\ No newline at end of file
diff --git a/test/low_rank_dynamics.jl b/test/low_rank_dynamics.jl
index 25008d82..8500aaae 100644
--- a/test/low_rank_dynamics.jl
+++ b/test/low_rank_dynamics.jl
@@ -66,7 +66,9 @@
         p0         = 0.,
         atol_inv   = 1e-6,
         adj_condition="variational",
-        Δt = 0.2, )
+        Δt = 0.2, 
+        progress = false
+    )
     lrsol = lr_mesolve(H, z, B, tl, c_ops; e_ops=e_ops, f_ops=(f_entropy,), opt=opt)
 
     # Test
diff --git a/test/progress_bar.jl b/test/progress_bar.jl
new file mode 100644
index 00000000..9ad0f1f1
--- /dev/null
+++ b/test/progress_bar.jl
@@ -0,0 +1,14 @@
+@testset "Progress Bar" begin
+    bar_width = 30
+    strLength = 67 + bar_width # including "\r" in the beginning of the string
+    prog = ProgressBar(bar_width, enable=true, bar_width=bar_width)
+    for p in 1:bar_width
+        output = sprint((t, s) -> next!(s, t), prog)
+
+        if p < bar_width
+            @test length(output) == strLength
+        else # the last output has an extra "\n" in the end
+            @test length(output) == strLength + 1
+        end
+    end
+end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 4f3d32fc..f783cc65 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -5,8 +5,8 @@ const GROUP = get(ENV, "GROUP", "All")
 
 const testdir = dirname(@__FILE__)
 
-# Put them in alphabetical order
-tests = [
+# Put core tests in alphabetical order
+core_tests = [
     "correlations_and_spectrum.jl",
     "dynamical_fock_dimension_mesolve.jl",
     "dynamical-shifted-fock.jl",
@@ -23,7 +23,7 @@ tests = [
 ]
 
 if (GROUP == "All") || (GROUP == "Core")
-    for test in tests
+    for test in core_tests
         include(joinpath(testdir, test))
     end
 end
\ No newline at end of file

From 48e674c984d42eb20d4585f087bc1c4515cde294 Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <y.t.d.huang@phys.ncku.edu.tw>
Date: Fri, 5 Apr 2024 18:09:18 +0800
Subject: [PATCH 2/2] add runtests for progress bar

---
 test/runtests.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/runtests.jl b/test/runtests.jl
index f783cc65..5841ad97 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -16,6 +16,7 @@ core_tests = [
     "low_rank_dynamics.jl",
     "negativity_and_partial_transpose.jl",
     "permutation.jl",
+    "progress_bar.jl",
     "quantum_objects.jl",
     "steady_state.jl",
     "time_evolution_and_partial_trace.jl",