1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189
|
import os
import pytest
from IPython.core.displaypub import CapturingDisplayPublisher
from IPython.core.interactiveshell import InteractiveShell
def run_in_ipython_shell(tmpdir, cells):
"""Run the given cells in an IPython shell and return the HTML output."""
InteractiveShell.clear_instance()
shell = InteractiveShell.instance(display_pub_class=CapturingDisplayPublisher)
prev_running_dir = os.getcwd()
try:
os.chdir(tmpdir)
for cell in cells:
shell.run_cell(cell)
finally:
os.chdir(prev_running_dir)
InteractiveShell.clear_instance()
try:
html = shell.display_pub.outputs[-1]["data"]["text/html"]
return html
except IndexError:
return None
@pytest.mark.filterwarnings("ignore")
class TestIPython:
def test_ipython_profiling(self, tmpdir):
"""Test that the IPython extension works."""
# GIVEN
code = [
"%load_ext memray",
"""
%%memray_flamegraph
x = "a" * 10000
""",
]
# WHEN
html = run_in_ipython_shell(tmpdir, code)
# THEN
assert html is not None
assert "<iframe" in html
assert "flamegraph.html" in html
def test_exception_while_ipython_profiling(self, tmpdir):
"""Test that the IPython extension works even if an exception
is raised."""
# GIVEN
code = [
"%load_ext memray",
"""
%%memray_flamegraph
x = "a" * 10000
1/0
""",
]
# WHEN
html = run_in_ipython_shell(tmpdir, code)
# THEN
assert html is not None
assert "<iframe" in html
assert "flamegraph.html" in html
def test_passing_help_argument(self, tmpdir, capsys):
# GIVEN
code = [
"%load_ext memray",
"""
%%memray_flamegraph -h
""",
]
# WHEN
html = run_in_ipython_shell(tmpdir, code)
# THEN
assert html is None
stdout, stderr = capsys.readouterr()
assert "show this help message" in stdout
assert "" == stderr
def test_passing_invalid_argument(self, tmpdir, capsys):
# GIVEN
code = [
"%load_ext memray",
"""
%%memray_flamegraph --oopsie
""",
]
# WHEN
html = run_in_ipython_shell(tmpdir, code)
# THEN
assert html is None
stdout, stderr = capsys.readouterr()
assert "" == stdout
assert "usage:" in stderr
def test_passing_valid_arguments(self, tmpdir, capsys):
# GIVEN
code = [
"%load_ext memray",
"""
%%memray_flamegraph --leaks --max-memory-records 5
x = "a" * 10000
""",
]
# WHEN
html = run_in_ipython_shell(tmpdir, code)
# THEN
assert html is not None
stdout, _ = capsys.readouterr()
assert "<iframe" in html
assert "flamegraph.html" in html
assert "Results saved to" in stdout
@pytest.mark.parametrize(
"args,title",
[
("--temporary-allocation-threshold=2", "flamegraph report"),
("--leaks", "flamegraph report (memory leaks)"),
("--temporal", "temporal flamegraph report"),
("--leaks --temporal", "temporal flamegraph report (memory leaks)"),
],
)
def test_report_title_by_report_type(self, tmpdir, capsys, args, title):
# GIVEN
code = [
"%load_ext memray",
f"""
%%memray_flamegraph {args}
x = "a" * 10000
""",
]
# WHEN
html = run_in_ipython_shell(tmpdir, code)
# THEN
assert html is not None
stdout, _ = capsys.readouterr()
assert "<iframe" in html
assert "flamegraph.html" in html
assert "Results saved to" in stdout
title_tag = f"<title>memray - {title}</title>"
assert title_tag in next(tmpdir.visit("**/flamegraph.html")).read()
def test_passing_temporal_and_temporary_allocations(self, tmpdir, capsys):
# GIVEN
code = [
"%load_ext memray",
"""
%%memray_flamegraph --temporal --temporary-allocation-threshold=2
""",
]
# WHEN
html = run_in_ipython_shell(tmpdir, code)
# THEN
assert html is None
stdout, stderr = capsys.readouterr()
assert "" == stdout
assert "Can't create a temporal flame graph of temporary allocations" in stderr
|