Add unit test for CLI options

* -a, --show-all (#64) * --color (#32: should work well with legacy watch)
wookayin · Aug 5, 2019 · 3cdaa54 · 3cdaa54
1 parent 2bf5377
commit 3cdaa54
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 6 deletions.
diff --git a/gpustat/cli.py b/gpustat/cli.py
@@ -118,7 +118,7 @@ def main(*argv):
         args.show_user = True
         args.show_pid = True
         args.show_fan_speed = True
-        args.show_power = ['draw,limit']
+        args.show_power = 'draw,limit'
     del args.show_all
 
     if args.interval is None:  # with default value

diff --git a/gpustat/test_gpustat.py b/gpustat/test_gpustat.py
@@ -166,6 +166,12 @@ def _MockedMem():
 """  # noqa: E501
 
 MOCK_EXPECTED_OUTPUT_FULL = """\
+[0] GeForce GTX TITAN 0 | 80°C,  16 %,  76 %,  125 / 250 W |  8000 / 12287 MB | user1:python/48448(4000M) user2:python/153223(4000M)
+[1] GeForce GTX TITAN 1 | 36°C,  53 %,   0 %,   ?? / 250 W |  9000 / 12189 MB | user1:torch/192453(3000M) user3:caffe/194826(6000M)
+[2] GeForce GTX TITAN 2 | 71°C, 100 %,  ?? %,  250 /  ?? W |     0 / 12189 MB | (Not Supported)
+"""  # noqa: E501
+
+MOCK_EXPECTED_OUTPUT_FULL_PROCESS = """\
 [0] GeForce GTX TITAN 0 | 80°C,  16 %,  76 %,  125 / 250 W |  8000 / 12287 MB | user1:python/48448(4000M) user2:python/153223(4000M)
  ├─  48448 (  85%,  257MB): python
  └─ 153223 (  15%,     0B): python
@@ -225,7 +231,7 @@ def test_new_query_mocked(self, N, Process, virtual_memory):
         unescaped = '\n'.join(unescaped.split('\n')[1:])
 
         self.maxDiff = 4096
-        self.assertEqual(unescaped, MOCK_EXPECTED_OUTPUT_FULL)
+        self.assertEqual(unescaped, MOCK_EXPECTED_OUTPUT_FULL_PROCESS)
 
     @mock.patch('psutil.virtual_memory')
     @mock.patch('psutil.Process')
@@ -290,16 +296,35 @@ def capture_output(*args):
                     )
             return f.getvalue()
 
+        def _remove_ansi_codes_and_header_line(s):
+            unescaped = remove_ansi_codes(s)
+            # remove first line (header)
+            unescaped = '\n'.join(unescaped.split('\n')[1:])
+            return unescaped
+
         s = capture_output('gpustat', )
-        unescaped = remove_ansi_codes(s)
-        # remove first line (header)
-        unescaped = '\n'.join(unescaped.split('\n')[1:])
         self.maxDiff = 4096
-        self.assertEqual(unescaped, MOCK_EXPECTED_OUTPUT_DEFAULT)
+        self.assertEqual(_remove_ansi_codes_and_header_line(s),
+                         MOCK_EXPECTED_OUTPUT_DEFAULT)
 
         s = capture_output('gpustat', '--no-header')
         self.assertIn("[0]", s.split('\n')[0])
 
+        s = capture_output('gpustat', '-a')  # --show-all
+        self.assertEqual(_remove_ansi_codes_and_header_line(s),
+                         MOCK_EXPECTED_OUTPUT_FULL)
+
+        s = capture_output('gpustat', '--color')
+        assert '\x0f' not in s, "Extra \\x0f found (see issue #32)"
+        self.assertEqual(_remove_ansi_codes_and_header_line(s),
+                         MOCK_EXPECTED_OUTPUT_DEFAULT)
+
+        s = capture_output('gpustat', '--no-color')
+        unescaped = remove_ansi_codes(s)
+        self.assertEqual(s, unescaped)   # should have no ansi code
+        self.assertEqual(_remove_ansi_codes_and_header_line(s),
+                         MOCK_EXPECTED_OUTPUT_DEFAULT)
+
     @mock.patch('psutil.virtual_memory')
     @mock.patch('psutil.Process')
     @mock.patch('gpustat.core.N')