test_progressbar.py 5.67 KB
Newer Older
Kai Chen's avatar
Kai Chen committed
1
# Copyright (c) Open-MMLab. All rights reserved.
lizz's avatar
lizz committed
2
import os
3
import time
Wenwei Zhang's avatar
Wenwei Zhang committed
4

lizz's avatar
lizz committed
5
6
7
8
9
try:
    from unittest.mock import patch
except ImportError:
    from mock import patch

lizz's avatar
lizz committed
10
11
12
13
try:
    from StringIO import StringIO
except ImportError:
    from io import StringIO
14

Wenwei Zhang's avatar
Wenwei Zhang committed
15
import mmcv  # isort:skip
16

lizz's avatar
lizz committed
17
18
19
20
21
22

def reset_string_io(io):
    io.truncate(0)
    io.seek(0)


23
24
class TestProgressBar(object):

lizz's avatar
lizz committed
25
26
    def test_start(self):
        out = StringIO()
27
28
        bar_width = 20
        # without total task num
lizz's avatar
lizz committed
29
30
31
32
33
34
        prog_bar = mmcv.ProgressBar(bar_width=bar_width, file=out)
        assert out.getvalue() == 'completed: 0, elapsed: 0s'
        reset_string_io(out)
        prog_bar = mmcv.ProgressBar(bar_width=bar_width, start=False, file=out)
        assert out.getvalue() == ''
        reset_string_io(out)
35
        prog_bar.start()
lizz's avatar
lizz committed
36
        assert out.getvalue() == 'completed: 0, elapsed: 0s'
37
        # with total task num
lizz's avatar
lizz committed
38
39
40
41
42
43
44
45
46
        reset_string_io(out)
        prog_bar = mmcv.ProgressBar(10, bar_width=bar_width, file=out)
        assert out.getvalue() == '[{}] 0/10, elapsed: 0s, ETA:'.format(
            ' ' * bar_width)
        reset_string_io(out)
        prog_bar = mmcv.ProgressBar(
            10, bar_width=bar_width, start=False, file=out)
        assert out.getvalue() == ''
        reset_string_io(out)
47
        prog_bar.start()
lizz's avatar
lizz committed
48
49
        assert out.getvalue() == '[{}] 0/10, elapsed: 0s, ETA:'.format(
            ' ' * bar_width)
50

lizz's avatar
lizz committed
51
52
    def test_update(self):
        out = StringIO()
53
54
        bar_width = 20
        # without total task num
lizz's avatar
lizz committed
55
        prog_bar = mmcv.ProgressBar(bar_width=bar_width, file=out)
56
        time.sleep(1)
lizz's avatar
lizz committed
57
        reset_string_io(out)
58
        prog_bar.update()
lizz's avatar
lizz committed
59
60
        assert out.getvalue() == 'completed: 1, elapsed: 1s, 1.0 tasks/s'
        reset_string_io(out)
61
        # with total task num
lizz's avatar
lizz committed
62
        prog_bar = mmcv.ProgressBar(10, bar_width=bar_width, file=out)
63
        time.sleep(1)
lizz's avatar
lizz committed
64
        reset_string_io(out)
65
        prog_bar.update()
lizz's avatar
lizz committed
66
67
68
        assert out.getvalue() == ('\r[{}] 1/10, 1.0 task/s, '
                                  'elapsed: 1s, ETA:     9s'.format('>' * 2 +
                                                                    ' ' * 18))
69

lizz's avatar
lizz committed
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
    def test_adaptive_length(self):
        with patch.dict('os.environ', {'COLUMNS': '80'}):
            out = StringIO()
            bar_width = 20
            prog_bar = mmcv.ProgressBar(10, bar_width=bar_width, file=out)
            time.sleep(1)
            reset_string_io(out)
            prog_bar.update()
            assert len(out.getvalue()) == 66

            os.environ['COLUMNS'] = '30'
            reset_string_io(out)
            prog_bar.update()
            assert len(out.getvalue()) == 48

            os.environ['COLUMNS'] = '60'
            reset_string_io(out)
            prog_bar.update()
            assert len(out.getvalue()) == 60

90
91
92
93
94
95

def sleep_1s(num):
    time.sleep(1)
    return num


lizz's avatar
lizz committed
96
97
98
99
100
101
102
103
def test_track_progress_list():
    out = StringIO()
    ret = mmcv.track_progress(sleep_1s, [1, 2, 3], bar_width=3, file=out)
    assert out.getvalue() == (
        '[   ] 0/3, elapsed: 0s, ETA:'
        '\r[>  ] 1/3, 1.0 task/s, elapsed: 1s, ETA:     2s'
        '\r[>> ] 2/3, 1.0 task/s, elapsed: 2s, ETA:     1s'
        '\r[>>>] 3/3, 1.0 task/s, elapsed: 3s, ETA:     0s\n')
104
105
106
    assert ret == [1, 2, 3]


lizz's avatar
lizz committed
107
108
def test_track_progress_iterator():
    out = StringIO()
109
    ret = mmcv.track_progress(
lizz's avatar
lizz committed
110
111
112
113
114
115
        sleep_1s, ((i for i in [1, 2, 3]), 3), bar_width=3, file=out)
    assert out.getvalue() == (
        '[   ] 0/3, elapsed: 0s, ETA:'
        '\r[>  ] 1/3, 1.0 task/s, elapsed: 1s, ETA:     2s'
        '\r[>> ] 2/3, 1.0 task/s, elapsed: 2s, ETA:     1s'
        '\r[>>>] 3/3, 1.0 task/s, elapsed: 3s, ETA:     0s\n')
116
117
118
    assert ret == [1, 2, 3]


lizz's avatar
lizz committed
119
120
def test_track_iter_progress():
    out = StringIO()
ZwwWayne's avatar
ZwwWayne committed
121
    ret = []
lizz's avatar
lizz committed
122
    for num in mmcv.track_iter_progress([1, 2, 3], bar_width=3, file=out):
ZwwWayne's avatar
ZwwWayne committed
123
        ret.append(sleep_1s(num))
lizz's avatar
lizz committed
124
125
126
127
128
    assert out.getvalue() == (
        '[   ] 0/3, elapsed: 0s, ETA:'
        '\r[>  ] 1/3, 1.0 task/s, elapsed: 1s, ETA:     2s'
        '\r[>> ] 2/3, 1.0 task/s, elapsed: 2s, ETA:     1s'
        '\r[>>>] 3/3, 1.0 task/s, elapsed: 3s, ETA:     0s\n')
ZwwWayne's avatar
ZwwWayne committed
129
130
131
    assert ret == [1, 2, 3]


lizz's avatar
lizz committed
132
133
def test_track_enum_progress():
    out = StringIO()
ZwwWayne's avatar
ZwwWayne committed
134
135
    ret = []
    count = []
lizz's avatar
lizz committed
136
137
    for i, num in enumerate(
            mmcv.track_iter_progress([1, 2, 3], bar_width=3, file=out)):
ZwwWayne's avatar
ZwwWayne committed
138
139
        ret.append(sleep_1s(num))
        count.append(i)
lizz's avatar
lizz committed
140
141
142
143
144
    assert out.getvalue() == (
        '[   ] 0/3, elapsed: 0s, ETA:'
        '\r[>  ] 1/3, 1.0 task/s, elapsed: 1s, ETA:     2s'
        '\r[>> ] 2/3, 1.0 task/s, elapsed: 2s, ETA:     1s'
        '\r[>>>] 3/3, 1.0 task/s, elapsed: 3s, ETA:     0s\n')
ZwwWayne's avatar
ZwwWayne committed
145
146
147
148
    assert ret == [1, 2, 3]
    assert count == [0, 1, 2]


lizz's avatar
lizz committed
149
150
def test_track_parallel_progress_list():
    out = StringIO()
151
    results = mmcv.track_parallel_progress(
lizz's avatar
lizz committed
152
153
154
155
156
157
158
        sleep_1s, [1, 2, 3, 4], 2, bar_width=4, file=out)
    assert out.getvalue() == (
        '[    ] 0/4, elapsed: 0s, ETA:'
        '\r[>   ] 1/4, 1.0 task/s, elapsed: 1s, ETA:     3s'
        '\r[>>  ] 2/4, 2.0 task/s, elapsed: 1s, ETA:     1s'
        '\r[>>> ] 3/4, 1.5 task/s, elapsed: 2s, ETA:     1s'
        '\r[>>>>] 4/4, 2.0 task/s, elapsed: 2s, ETA:     0s\n')
159
160
161
    assert results == [1, 2, 3, 4]


lizz's avatar
lizz committed
162
163
def test_track_parallel_progress_iterator():
    out = StringIO()
164
    results = mmcv.track_parallel_progress(
lizz's avatar
lizz committed
165
166
167
168
169
170
171
        sleep_1s, ((i for i in [1, 2, 3, 4]), 4), 2, bar_width=4, file=out)
    assert out.getvalue() == (
        '[    ] 0/4, elapsed: 0s, ETA:'
        '\r[>   ] 1/4, 1.0 task/s, elapsed: 1s, ETA:     3s'
        '\r[>>  ] 2/4, 2.0 task/s, elapsed: 1s, ETA:     1s'
        '\r[>>> ] 3/4, 1.5 task/s, elapsed: 2s, ETA:     1s'
        '\r[>>>>] 4/4, 2.0 task/s, elapsed: 2s, ETA:     0s\n')
172
    assert results == [1, 2, 3, 4]