test_progressbar.py 5.81 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
Kai Chen's avatar
Kai Chen committed
3
import sys
4
import time
Wenwei Zhang's avatar
Wenwei Zhang committed
5

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

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

Wenwei Zhang's avatar
Wenwei Zhang committed
16
import pytest  # isort:skip
Kai Chen's avatar
Kai Chen committed
17

Wenwei Zhang's avatar
Wenwei Zhang committed
18
import mmcv  # isort:skip
19

lizz's avatar
lizz committed
20
21
22
23
24
25

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


Kai Chen's avatar
Kai Chen committed
26
27
28
if sys.version_info[0] == 2:
    pytest.skip('skipping tests for python 2', allow_module_level=True)

29
30
31

class TestProgressBar(object):

lizz's avatar
lizz committed
32
33
    def test_start(self):
        out = StringIO()
34
35
        bar_width = 20
        # without total task num
lizz's avatar
lizz committed
36
37
38
39
40
41
        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)
42
        prog_bar.start()
lizz's avatar
lizz committed
43
        assert out.getvalue() == 'completed: 0, elapsed: 0s'
44
        # with total task num
lizz's avatar
lizz committed
45
46
47
48
49
50
51
52
53
        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)
54
        prog_bar.start()
lizz's avatar
lizz committed
55
56
        assert out.getvalue() == '[{}] 0/10, elapsed: 0s, ETA:'.format(
            ' ' * bar_width)
57

lizz's avatar
lizz committed
58
59
    def test_update(self):
        out = StringIO()
60
61
        bar_width = 20
        # without total task num
lizz's avatar
lizz committed
62
        prog_bar = mmcv.ProgressBar(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
        assert out.getvalue() == 'completed: 1, elapsed: 1s, 1.0 tasks/s'
        reset_string_io(out)
68
        # with total task num
lizz's avatar
lizz committed
69
        prog_bar = mmcv.ProgressBar(10, bar_width=bar_width, file=out)
70
        time.sleep(1)
lizz's avatar
lizz committed
71
        reset_string_io(out)
72
        prog_bar.update()
lizz's avatar
lizz committed
73
74
75
        assert out.getvalue() == ('\r[{}] 1/10, 1.0 task/s, '
                                  'elapsed: 1s, ETA:     9s'.format('>' * 2 +
                                                                    ' ' * 18))
76

lizz's avatar
lizz committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
    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

97
98
99
100
101
102

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


lizz's avatar
lizz committed
103
104
105
106
107
108
109
110
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')
111
112
113
    assert ret == [1, 2, 3]


lizz's avatar
lizz committed
114
115
def test_track_progress_iterator():
    out = StringIO()
116
    ret = mmcv.track_progress(
lizz's avatar
lizz committed
117
118
119
120
121
122
        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')
123
124
125
    assert ret == [1, 2, 3]


lizz's avatar
lizz committed
126
127
def test_track_iter_progress():
    out = StringIO()
ZwwWayne's avatar
ZwwWayne committed
128
    ret = []
lizz's avatar
lizz committed
129
    for num in mmcv.track_iter_progress([1, 2, 3], bar_width=3, file=out):
ZwwWayne's avatar
ZwwWayne committed
130
        ret.append(sleep_1s(num))
lizz's avatar
lizz committed
131
132
133
134
135
    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
136
137
138
    assert ret == [1, 2, 3]


lizz's avatar
lizz committed
139
140
def test_track_enum_progress():
    out = StringIO()
ZwwWayne's avatar
ZwwWayne committed
141
142
    ret = []
    count = []
lizz's avatar
lizz committed
143
144
    for i, num in enumerate(
            mmcv.track_iter_progress([1, 2, 3], bar_width=3, file=out)):
ZwwWayne's avatar
ZwwWayne committed
145
146
        ret.append(sleep_1s(num))
        count.append(i)
lizz's avatar
lizz committed
147
148
149
150
151
    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
152
153
154
155
    assert ret == [1, 2, 3]
    assert count == [0, 1, 2]


lizz's avatar
lizz committed
156
157
def test_track_parallel_progress_list():
    out = StringIO()
158
    results = mmcv.track_parallel_progress(
lizz's avatar
lizz committed
159
160
161
162
163
164
165
        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')
166
167
168
    assert results == [1, 2, 3, 4]


lizz's avatar
lizz committed
169
170
def test_track_parallel_progress_iterator():
    out = StringIO()
171
    results = mmcv.track_parallel_progress(
lizz's avatar
lizz committed
172
173
174
175
176
177
178
        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')
179
    assert results == [1, 2, 3, 4]