add tests for write_images() and to_images()

Author: emilykl

tests/test_optional/test_kaleido/test_kaleido.py

@@ -3,6 +3,8 @@
 import tempfile
 from contextlib import redirect_stdout
 import base64
+import unittest
+from unittest import mock
 
 from pdfrw import PdfReader
 from PIL import Image
@@ -121,3 +123,131 @@ def test_bytesio():
     bio_bytes = bio.read()
     to_image_bytes = pio.to_image(fig, format="jpg", engine="kaleido", validate=False)
     assert bio_bytes == to_image_bytes
+
+
+@mock.patch("plotly.io._kaleido.to_image")
+def test_to_images_single(mock_to_image):
+    """Test to_images with a single figure"""
+    pio.to_images(
+        fig,
+        format="png",
+        width=800,
+        height=600,
+        scale=2,
+        validate=True,
+    )
+
+    # Verify that to_image was called once with the correct arguments
+    expected_calls = [
+        mock.call(
+            fig,
+            format="png",
+            width=800,
+            height=600,
+            scale=2,
+            validate=True,
+        )
+    ]
+    mock_to_image.assert_has_calls(expected_calls, any_order=False)
+    assert mock_to_image.call_count == 1
+
+
+@mock.patch("plotly.io._kaleido.to_image")
+def test_to_images_multiple(mock_to_image):
+    """Test to_images with lists"""
+    fig1 = {"data": [], "layout": {"title": {"text": "figure 1"}}}
+    fig2 = {"data": [], "layout": {"title": {"text": "figure 2"}}}
+    pio.to_images(
+        [fig1, fig2],
+        "png",
+        width=[800, 400],
+        height=600,
+        scale=[1, 2],
+        validate=True,
+    )
+
+    # Verify that to_image was called with the correct arguments in the correct order
+    expected_calls = [
+        mock.call(
+            fig1,
+            "png",
+            width=800,
+            height=600,
+            scale=1,
+            validate=True,
+        ),
+        mock.call(
+            fig2,
+            "png",
+            width=400,
+            height=600,
+            scale=2,
+            validate=True,
+        ),
+    ]
+    mock_to_image.assert_has_calls(expected_calls, any_order=False)
+    assert mock_to_image.call_count == 2
+
+
+@mock.patch("plotly.io._kaleido.write_image")
+def test_write_images_single(mock_write_image):
+    """Test write_images with only single arguments"""
+    pio.write_images(
+        fig,
+        "output.png",
+        format="png",
+        width=800,
+        height=600,
+        scale=2,
+    )
+
+    # Verify that write_image was called once with the correct arguments
+    expected_calls = [
+        mock.call(
+            fig,
+            "output.png",
+            format="png",
+            width=800,
+            height=600,
+            scale=2,
+        )
+    ]
+    mock_write_image.assert_has_calls(expected_calls, any_order=False)
+    assert mock_write_image.call_count == 1
+
+
+@mock.patch("plotly.io._kaleido.write_image")
+def test_write_images_multiple(mock_write_image):
+    """Test write_images with list arguments"""
+    fig1 = {"data": [], "layout": {"title": {"text": "figure 1"}}}
+    fig2 = {"data": [], "layout": {"title": {"text": "figure 2"}}}
+    pio.write_images(
+        [fig1, fig2],
+        ["output1.png", "output2.jpg"],
+        format=["png", "jpeg"],
+        width=800,
+        height=[600, 400],
+        scale=2,
+    )
+
+    # Verify that write_image was called with the correct arguments in the correct order
+    expected_calls = [
+        mock.call(
+            fig1,
+            "output1.png",
+            format="png",
+            width=800,
+            height=600,
+            scale=2,
+        ),
+        mock.call(
+            fig2,
+            "output2.jpg",
+            format="jpeg",
+            width=800,
+            height=400,
+            scale=2,
+        ),
+    ]
+    mock_write_image.assert_has_calls(expected_calls, any_order=False)
+    assert mock_write_image.call_count == 2