Skip to content

Commit 6c25793

Browse files
committed
add test to prevent regression
1 parent 7250704 commit 6c25793

File tree

1 file changed

+45
-2
lines changed

1 file changed

+45
-2
lines changed

tests/test_optional/test_kaleido/test_kaleido.py

Lines changed: 45 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,18 @@
1+
import base64
2+
from contextlib import redirect_stdout
13
from io import BytesIO, StringIO
24
from pathlib import Path
35
import tempfile
4-
from contextlib import redirect_stdout
5-
import base64
6+
from unittest.mock import patch
67

78
from pdfrw import PdfReader
89
from PIL import Image
10+
import plotly.graph_objects as go
911
import plotly.io as pio
1012
from plotly.io.kaleido import kaleido_available, kaleido_major
1113
import pytest
1214

15+
1316
fig = {"data": [], "layout": {"title": {"text": "figure title"}}}
1417

1518

@@ -160,3 +163,43 @@ def test_defaults():
160163
finally:
161164
pio.defaults.default_format = "png"
162165
assert pio.defaults.default_format == "png"
166+
167+
168+
def test_fig_write_image():
169+
"""Test that fig.write_image() calls the correct underlying Kaleido function."""
170+
171+
test_fig = go.Figure(fig)
172+
test_image_bytes = b"mock image data"
173+
174+
if kaleido_major() > 0:
175+
patch_funcname = "plotly.io._kaleido.kaleido.calc_fig_sync"
176+
else:
177+
patch_funcname = "plotly.io._kaleido.scope.transform"
178+
179+
with patch(patch_funcname, return_value=test_image_bytes) as mock_calc_fig:
180+
test_fig.write_image("test_path.png")
181+
182+
# Verify patched function was called once with fig dict as first argument
183+
mock_calc_fig.assert_called_once()
184+
args, _ = mock_calc_fig.call_args
185+
assert args[0] == test_fig.to_dict()
186+
187+
188+
def test_fig_to_image():
189+
"""Test that fig.to_image() calls the correct underlying Kaleido function."""
190+
191+
test_fig = go.Figure(fig)
192+
test_image_bytes = b"mock image data"
193+
194+
if kaleido_major() > 0:
195+
patch_funcname = "plotly.io._kaleido.kaleido.calc_fig_sync"
196+
else:
197+
patch_funcname = "plotly.io._kaleido.scope.transform"
198+
199+
with patch(patch_funcname, return_value=test_image_bytes) as mock_calc_fig:
200+
test_fig.to_image()
201+
202+
# Verify patched function was called once with fig dict as first argument
203+
mock_calc_fig.assert_called_once()
204+
args, _ = mock_calc_fig.call_args
205+
assert args[0] == test_fig.to_dict()

0 commit comments

Comments
 (0)