Skip to content

Commit 9cf495d

Browse files
committed
fix(tests): use ExitStack for Python 3.9 compatibility
1 parent 8ba6845 commit 9cf495d

1 file changed

Lines changed: 80 additions & 46 deletions

File tree

datalab_kernel/tests/unit/test_backend_selection.py

Lines changed: 80 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
from __future__ import annotations
1313

14+
import contextlib
1415
import os
1516
from unittest import mock
1617

@@ -65,79 +66,112 @@ def test_invalid_backend_raises(self):
6566

6667
def test_auto_prefers_plotly(self):
6768
"""Auto-detect prefers plotly when both are available."""
68-
with (
69-
mock.patch("datalab_kernel.plotter.plotly_available", return_value=True),
70-
mock.patch(
71-
"datalab_kernel.plotter.matplotlib_available", return_value=True
72-
),
73-
):
69+
with contextlib.ExitStack() as stack:
70+
stack.enter_context(
71+
mock.patch("datalab_kernel.plotter.plotly_available", return_value=True)
72+
)
73+
stack.enter_context(
74+
mock.patch(
75+
"datalab_kernel.plotter.matplotlib_available", return_value=True
76+
)
77+
)
7478
assert resolve_backend(None) == BACKEND_PLOTLY
7579

7680
def test_auto_falls_back_to_matplotlib(self):
7781
"""Auto-detect falls back to matplotlib when plotly is missing."""
78-
with (
79-
mock.patch("datalab_kernel.plotter.plotly_available", return_value=False),
80-
mock.patch(
81-
"datalab_kernel.plotter.matplotlib_available", return_value=True
82-
),
83-
):
82+
with contextlib.ExitStack() as stack:
83+
stack.enter_context(
84+
mock.patch(
85+
"datalab_kernel.plotter.plotly_available", return_value=False
86+
)
87+
)
88+
stack.enter_context(
89+
mock.patch(
90+
"datalab_kernel.plotter.matplotlib_available", return_value=True
91+
)
92+
)
8493
assert resolve_backend(None) == BACKEND_MATPLOTLIB
8594

8695
def test_auto_plotly_only(self):
8796
"""Auto-detect returns plotly when only plotly is available."""
88-
with (
89-
mock.patch("datalab_kernel.plotter.plotly_available", return_value=True),
90-
mock.patch(
91-
"datalab_kernel.plotter.matplotlib_available", return_value=False
92-
),
93-
):
97+
with contextlib.ExitStack() as stack:
98+
stack.enter_context(
99+
mock.patch(
100+
"datalab_kernel.plotter.plotly_available", return_value=True
101+
)
102+
)
103+
stack.enter_context(
104+
mock.patch(
105+
"datalab_kernel.plotter.matplotlib_available", return_value=False
106+
)
107+
)
94108
assert resolve_backend(None) == BACKEND_PLOTLY
95109

96110
def test_neither_raises(self):
97111
"""ImportError when neither backend is available."""
98-
with (
99-
mock.patch("datalab_kernel.plotter.plotly_available", return_value=False),
100-
mock.patch(
101-
"datalab_kernel.plotter.matplotlib_available", return_value=False
102-
),
103-
pytest.raises(ImportError, match="Neither plotly nor matplotlib"),
104-
):
105-
resolve_backend(None)
112+
with contextlib.ExitStack() as stack:
113+
stack.enter_context(
114+
mock.patch(
115+
"datalab_kernel.plotter.plotly_available", return_value=False
116+
)
117+
)
118+
stack.enter_context(
119+
mock.patch(
120+
"datalab_kernel.plotter.matplotlib_available", return_value=False
121+
)
122+
)
123+
with pytest.raises(ImportError, match="Neither plotly nor matplotlib"):
124+
resolve_backend(None)
106125

107126
def test_fallback_with_warning(self):
108127
"""Falls back to the other backend with a warning."""
109-
with (
110-
mock.patch("datalab_kernel.plotter.plotly_available", return_value=False),
111-
mock.patch(
112-
"datalab_kernel.plotter.matplotlib_available", return_value=True
113-
),
114-
):
128+
with contextlib.ExitStack() as stack:
129+
stack.enter_context(
130+
mock.patch(
131+
"datalab_kernel.plotter.plotly_available", return_value=False
132+
)
133+
)
134+
stack.enter_context(
135+
mock.patch(
136+
"datalab_kernel.plotter.matplotlib_available", return_value=True
137+
)
138+
)
115139
with pytest.warns(UserWarning, match="not installed.*Falling back"):
116140
result = resolve_backend("plotly")
117141
assert result == BACKEND_MATPLOTLIB
118142

119143
def test_fallback_other_direction(self):
120144
"""Falls back to plotly when matplotlib is requested but missing."""
121-
with (
122-
mock.patch(
123-
"datalab_kernel.plotter.matplotlib_available", return_value=False
124-
),
125-
mock.patch("datalab_kernel.plotter.plotly_available", return_value=True),
126-
):
145+
with contextlib.ExitStack() as stack:
146+
stack.enter_context(
147+
mock.patch(
148+
"datalab_kernel.plotter.matplotlib_available", return_value=False
149+
)
150+
)
151+
stack.enter_context(
152+
mock.patch(
153+
"datalab_kernel.plotter.plotly_available", return_value=True
154+
)
155+
)
127156
with pytest.warns(UserWarning, match="not installed.*Falling back"):
128157
result = resolve_backend("matplotlib")
129158
assert result == BACKEND_PLOTLY
130159

131160
def test_both_missing_with_explicit_raises(self):
132161
"""ImportError when explicit backend and fallback are both missing."""
133-
with (
134-
mock.patch("datalab_kernel.plotter.plotly_available", return_value=False),
135-
mock.patch(
136-
"datalab_kernel.plotter.matplotlib_available", return_value=False
137-
),
138-
pytest.raises(ImportError, match="Neither plotly nor matplotlib"),
139-
):
140-
resolve_backend("plotly")
162+
with contextlib.ExitStack() as stack:
163+
stack.enter_context(
164+
mock.patch(
165+
"datalab_kernel.plotter.plotly_available", return_value=False
166+
)
167+
)
168+
stack.enter_context(
169+
mock.patch(
170+
"datalab_kernel.plotter.matplotlib_available", return_value=False
171+
)
172+
)
173+
with pytest.raises(ImportError, match="Neither plotly nor matplotlib"):
174+
resolve_backend("plotly")
141175

142176

143177
# ============================================================================

0 commit comments

Comments
 (0)