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

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

openmc/universe.py

@@ -370,8 +370,8 @@ class Universe(UniverseBase):
 
         Returns
         -------
-        matplotlib.image.AxesImage
-            Resulting image
+        matplotlib.axes.Axes
+            Axes containing resulting image
 
         """
         import matplotlib.image as mpimg
@@ -483,7 +483,7 @@ class Universe(UniverseBase):
             # 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 @@ class Universe(UniverseBase):
                 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.

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,7 +80,7 @@ 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,
@@ -85,6 +88,17 @@ def test_plot(run_in_tmpdir, sphere_model):
             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):