|
8 | 8 | from test.support.script_helper import (assert_python_ok, assert_python_failure, |
9 | 9 | interpreter_requires_environment) |
10 | 10 | from test import support |
11 | | -from test.support import force_not_colorized, warnings_helper |
| 11 | +from test.support import (force_not_colorized, |
| 12 | + force_not_colorized_test_class, warnings_helper) |
12 | 13 | from test.support import os_helper |
13 | 14 | from test.support import threading_helper |
14 | 15 |
|
@@ -1000,6 +1001,149 @@ def test_pymem_alloc0(self): |
1000 | 1001 | assert_python_ok('-X', 'tracemalloc', '-c', code) |
1001 | 1002 |
|
1002 | 1003 |
|
| 1004 | +@force_not_colorized_test_class |
| 1005 | +class TestModuleCommandLine(unittest.TestCase): |
| 1006 | + script = textwrap.dedent(""" |
| 1007 | + import sys |
| 1008 | + def alloc(): |
| 1009 | + return [bytes(10_000) for _ in range(100)] |
| 1010 | + data = alloc() |
| 1011 | + print(f'name={__name__} argv={sys.argv[1:]}') |
| 1012 | + """) |
| 1013 | + |
| 1014 | + def setUp(self): |
| 1015 | + self.tmpdir = self.enterContext(os_helper.temp_dir()) |
| 1016 | + self.script_path = self.write_script('script.py', self.script) |
| 1017 | + |
| 1018 | + def write_script(self, name, source): |
| 1019 | + path = os.path.join(self.tmpdir, name) |
| 1020 | + with open(path, 'w', encoding='utf-8') as fp: |
| 1021 | + fp.write(source) |
| 1022 | + return path |
| 1023 | + |
| 1024 | + def run_cli(self, *args, **kwargs): |
| 1025 | + res = assert_python_ok('-m', 'tracemalloc', *args, **kwargs) |
| 1026 | + return res.out.decode() |
| 1027 | + |
| 1028 | + def run_cli_failure(self, *args): |
| 1029 | + res = assert_python_failure('-m', 'tracemalloc', *args) |
| 1030 | + return res.err.decode() |
| 1031 | + |
| 1032 | + def test_script(self): |
| 1033 | + out = self.run_cli(self.script_path, 'a', '-b', '--c') |
| 1034 | + self.assertIn("name=__main__ argv=['a', '-b', '--c']", out) |
| 1035 | + self.assertIn('Top 10 lineno allocations still alive at exit:', out) |
| 1036 | + self.assertRegex(out, r'script\.py:4: size=\d+ KiB, count=\d+') |
| 1037 | + self.assertIn('Total size: ', out) |
| 1038 | + self.assertIn('Peak size: ', out) |
| 1039 | + |
| 1040 | + def test_globals_alive(self): |
| 1041 | + # Memory referenced by globals is alive at exit, even without a |
| 1042 | + # reference cycle between a function and the globals |
| 1043 | + script = self.write_script( |
| 1044 | + 'nofunc.py', 'data = [bytes(10_000) for _ in range(100)]\n') |
| 1045 | + out = self.run_cli(script) |
| 1046 | + self.assertRegex(out, r'nofunc\.py:1: size=\d+ KiB, count=\d+') |
| 1047 | + |
| 1048 | + def test_module(self): |
| 1049 | + out = self.run_cli('-m', 'script', 'arg', __cwd=self.tmpdir, |
| 1050 | + __isolated=False) |
| 1051 | + self.assertIn("name=__main__ argv=['arg']", out) |
| 1052 | + self.assertRegex(out, r'script\.py:4: size=\d+ KiB') |
| 1053 | + |
| 1054 | + def test_script_sys_path(self): |
| 1055 | + self.write_script('helper.py', 'print("helper imported")\n') |
| 1056 | + script = self.write_script('importer.py', 'import helper\n') |
| 1057 | + out = self.run_cli(script, __isolated=False) |
| 1058 | + self.assertIn('helper imported', out) |
| 1059 | + |
| 1060 | + def test_limit(self): |
| 1061 | + out = self.run_cli('-l', '1', self.script_path) |
| 1062 | + lines = out.splitlines() |
| 1063 | + index = lines.index('Top 1 lineno allocations still alive at exit:') |
| 1064 | + self.assertRegex(lines[index + 1], r'script\.py:4: ') |
| 1065 | + self.assertRegex(lines[index + 2], r'^\d+ other: size=') |
| 1066 | + |
| 1067 | + def test_key_type_filename(self): |
| 1068 | + out = self.run_cli('-t', 'filename', self.script_path) |
| 1069 | + self.assertRegex(out, r'script\.py:0: size=\d+ KiB') |
| 1070 | + |
| 1071 | + def test_key_type_traceback(self): |
| 1072 | + out = self.run_cli('-t', 'traceback', '-n', '5', '-l', '1', |
| 1073 | + self.script_path) |
| 1074 | + self.assertIn('Top 1 traceback allocations still alive at exit:', out) |
| 1075 | + self.assertIn(' data = alloc()', out) |
| 1076 | + self.assertIn(' return [bytes(10_000) for _ in range(100)]', out) |
| 1077 | + # Frames of the command line interface are not displayed |
| 1078 | + self.assertNotIn('runpy', out) |
| 1079 | + self.assertNotIn('tracemalloc.py', out) |
| 1080 | + |
| 1081 | + def test_cumulative(self): |
| 1082 | + out = self.run_cli('--cumulative', '-n', '5', '-l', '2', |
| 1083 | + self.script_path) |
| 1084 | + self.assertRegex(out, r'script\.py:5: size=\d+ KiB') |
| 1085 | + self.assertRegex(out, r'script\.py:4: size=\d+ KiB') |
| 1086 | + self.assertNotIn('runpy', out) |
| 1087 | + self.assertNotIn(' other: ', out) |
| 1088 | + |
| 1089 | + def test_outfile(self): |
| 1090 | + outfile = os.path.join(self.tmpdir, 'snapshot.pickle') |
| 1091 | + out = self.run_cli('-n', '3', '-o', outfile, self.script_path) |
| 1092 | + self.assertNotIn('Total size', out) |
| 1093 | + snapshot = tracemalloc.Snapshot.load(outfile) |
| 1094 | + self.assertEqual(snapshot.traceback_limit, 3) |
| 1095 | + top = snapshot.statistics('lineno')[0] |
| 1096 | + self.assertEqual(top.traceback[0].filename, self.script_path) |
| 1097 | + self.assertEqual(top.traceback[0].lineno, 4) |
| 1098 | + |
| 1099 | + def test_already_tracing(self): |
| 1100 | + # -n overrides the number of frames set by -X tracemalloc |
| 1101 | + outfile = os.path.join(self.tmpdir, 'snapshot.pickle') |
| 1102 | + self.run_cli('-n', '3', '-o', outfile, self.script_path, |
| 1103 | + __isolated=False, PYTHONTRACEMALLOC='1') |
| 1104 | + snapshot = tracemalloc.Snapshot.load(outfile) |
| 1105 | + self.assertEqual(snapshot.traceback_limit, 3) |
| 1106 | + |
| 1107 | + def test_exit_code(self): |
| 1108 | + script = self.write_script('exit.py', 'import sys\nsys.exit(3)\n') |
| 1109 | + res = assert_python_failure('-m', 'tracemalloc', script) |
| 1110 | + self.assertEqual(res.rc, 3) |
| 1111 | + self.assertIn(b'Total size: ', res.out) |
| 1112 | + |
| 1113 | + def test_exception(self): |
| 1114 | + script = self.write_script('error.py', '1/0\n') |
| 1115 | + res = assert_python_failure('-m', 'tracemalloc', script) |
| 1116 | + self.assertIn(b'ZeroDivisionError', res.err) |
| 1117 | + self.assertIn(b'Total size: ', res.out) |
| 1118 | + |
| 1119 | + def test_tracing_stopped(self): |
| 1120 | + # The traced program can stop tracing |
| 1121 | + script = self.write_script( |
| 1122 | + 'stop.py', 'import sys, tracemalloc\n' |
| 1123 | + 'tracemalloc.stop()\n' |
| 1124 | + 'sys.exit(3)\n') |
| 1125 | + res = assert_python_failure('-m', 'tracemalloc', script) |
| 1126 | + self.assertEqual(res.rc, 3) |
| 1127 | + self.assertEqual(res.out, b'') |
| 1128 | + self.assertEqual(res.err, b'') |
| 1129 | + |
| 1130 | + def test_errors(self): |
| 1131 | + err = self.run_cli_failure() |
| 1132 | + self.assertIn('the following arguments are required: target', err) |
| 1133 | + err = self.run_cli_failure(os.path.join(self.tmpdir, 'missing.py')) |
| 1134 | + self.assertIn('cannot find script', err) |
| 1135 | + err = self.run_cli_failure('-t', 'traceback', '--cumulative', |
| 1136 | + self.script_path) |
| 1137 | + self.assertIn("--cumulative cannot be used with key type", err) |
| 1138 | + err = self.run_cli_failure('-l', '0', self.script_path) |
| 1139 | + self.assertIn('--limit must be a positive integer', err) |
| 1140 | + for nframe in INVALID_NFRAME: |
| 1141 | + with self.subTest(nframe=nframe): |
| 1142 | + err = self.run_cli_failure('-n', str(nframe), |
| 1143 | + self.script_path) |
| 1144 | + self.assertIn('the number of frames must be in range', err) |
| 1145 | + |
| 1146 | + |
1003 | 1147 | @unittest.skipIf(_testcapi is None, 'need _testcapi') |
1004 | 1148 | class TestCAPI(unittest.TestCase): |
1005 | 1149 | maxDiff = 80 * 20 |
|
0 commit comments