added universe plot tests to catch color=None and axis values (#2637)

Co-authored-by: Paul Romano <[email protected]>
openmc-dev · Aug 19, 2023 · 0964024 · 0964024
1 parent b62b62a
commit 0964024
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 7 deletions.
diff --git a/openmc/universe.py b/openmc/universe.py
@@ -370,8 +370,8 @@ def plot(self, origin=None, width=None, pixels=40000,
 
         Returns
         -------
-        matplotlib.image.AxesImage
-            Resulting image
+        matplotlib.axes.Axes
+            Axes containing resulting image
 
         """
         import matplotlib.image as mpimg
@@ -483,7 +483,7 @@ def plot(self, origin=None, width=None, pixels=40000,
             # add legend showing which colors represent which material
             # or cell if that was requested
             if legend:
-                if plot.colors is None:
+                if plot.colors == {}:
                     raise ValueError("Must pass 'colors' dictionary if you "
                                      "are adding a legend via legend=True.")
 
@@ -522,7 +522,8 @@ def plot(self, origin=None, width=None, pixels=40000,
                 axes.legend(handles=patches, **legend_kwargs)
 
             # Plot image and return the axes
-            return axes.imshow(img, extent=(x_min, x_max, y_min, y_max), **kwargs)
+            axes.imshow(img, extent=(x_min, x_max, y_min, y_max), **kwargs)
+            return axes
 
     def add_cell(self, cell):
         """Add a cell to the universe.

diff --git a/tests/unit_tests/test_universe.py b/tests/unit_tests/test_universe.py
@@ -61,14 +61,17 @@ def test_plot(run_in_tmpdir, sphere_model):
     }
 
     for basis in ('xy', 'yz', 'xz'):
-        pincell.geometry.root_universe.plot(
+        plot = pincell.geometry.root_universe.plot(
             colors=mat_colors,
             color_by="material",
             legend=True,
             pixels=(10, 10),
             basis=basis,
-            outline=True
+            outline=True,
+            axis_units='m'
         )
+        assert plot.xaxis.get_label().get_text() == f'{basis[0]} [m]'
+        assert plot.yaxis.get_label().get_text() == f'{basis[1]} [m]'
 
     # model with no inf values in bounding box
     m = sphere_model.materials[0]
@@ -77,14 +80,25 @@ def test_plot(run_in_tmpdir, sphere_model):
     colors = {m: 'limegreen'}
 
     for basis in ('xy', 'yz', 'xz'):
-        univ.plot(
+        plot = univ.plot(
             colors=colors,
             color_by="cell",
             legend=False,
             pixels=100,
             basis=basis,
             outline=False
         )
+        assert plot.xaxis.get_label().get_text() == f'{basis[0]} [cm]'
+        assert plot.yaxis.get_label().get_text() == f'{basis[1]} [cm]'
+
+    msg = "Must pass 'colors' dictionary if you are adding a legend via legend=True."
+    # This plot call should fail as legend is True but colors is None
+    with pytest.raises(ValueError, match=msg):
+        univ.plot(
+            color_by="cell",
+            legend=True,
+            pixels=100,
+        )
 
 
 def test_get_nuclides(uo2):