test_gpustat.py 文件源码

def test_args_endtoend(self, N, Process):
        """
        End-to-end testing given command line args.
        """
        _configure_mock(N, Process)

        def capture_output(*args):
            f = StringIO()
            import contextlib

            with contextlib.redirect_stdout(f):  # requires python 3.4+
                try:
                    gpustat.main(*args)
                except SystemExit:
                    raise AssertionError("Argparse failed (see above error message)")
            return f.getvalue()

        s = capture_output('gpustat', )
        unescaped = remove_ansi_codes(s)
        unescaped = '\n'.join(unescaped.split('\n')[1:])   # remove first line (header)
        self.maxDiff = 4096
        self.assertEqual(unescaped, MOCK_EXPECTED_OUTPUT_DEFAULT)

        s = capture_output('gpustat', '--no-header')
        self.assertIn("[0]", s.split('\n')[0])