Tests for go.Figure.set_subplots and figure argument to make_subplots

Author: nicholas-esterer

packages/python/plotly/plotly/tests/test_core/test_graph_objs/test_figure.py

@@ -1,6 +1,7 @@
 from __future__ import absolute_import
 
 import plotly.graph_objects as go
+from plotly.subplots import make_subplots
 from plotly.tests.utils import TestCaseNoTemplate
 
 
@@ -199,3 +200,20 @@ def test_update_overwrite_data(self):
             fig.to_plotly_json()["data"],
             [{"type": "scatter", "y": [1, 3, 2], "line": {"color": "yellow"}}],
         )
+
+    def test_set_subplots(self):
+        # Test that it works the same as make_subplots for a simple call
+        fig0 = go.Figure()
+        fig0_sp = make_subplots(2, 2)
+        fig0.set_subplots(2, 2)
+        assert fig0.layout == fig0_sp.layout
+        # Test that it accepts the same arguments as make_subplots
+        fig1 = go.Figure()
+        fig1.set_subplots(rows=2, cols=2, horizontal_spacing=0.25, vertical_spacing=0.1)
+        fig1_sp = make_subplots(
+            rows=2, cols=2, horizontal_spacing=0.25, vertical_spacing=0.1
+        )
+        assert fig1.layout == fig1_sp.layout
+        # Test that calling on a figure that already has subplots throws an error.
+        with self.assertRaisesRegex(ValueError, "^This figure already has subplots\.$"):
+            fig1.set_subplots(2, 3)

packages/python/plotly/plotly/tests/test_core/test_subplots/test_make_subplots.py

@@ -1924,3 +1924,13 @@ def test_secondary_y_subplots(self):
         expected.update_traces(uid=None)
 
         self.assertEqual(fig.to_plotly_json(), expected.to_plotly_json())
+
+    def test_if_passed_figure(self):
+        # assert it returns the same figure it was passed
+        fig = Figure()
+        figsp = subplots.make_subplots(2, 2, figure=fig)
+        assert id(fig) == id(figsp)
+        # assert the layout is the same when it returns its own figure
+        fig2sp = subplots.make_subplots(2, 2)
+        assert id(fig2sp) != id(figsp)
+        assert fig2sp.layout == figsp.layout