From 7f16f46beddff4d6090a154dc5270e35d31ce5a3 Mon Sep 17 00:00:00 2001
From: shimwell <mail@jshimwell.com>
Date: Thu, 2 Mar 2023 23:52:00 +0000
Subject: [PATCH] material and plot_xs tests for plotter.py

---
 tests/unit_tests/test_plotter.py | 46 ++++++++++++++++++++++++++------
 1 file changed, 38 insertions(+), 8 deletions(-)

diff --git a/tests/unit_tests/test_plotter.py b/tests/unit_tests/test_plotter.py
index be420c3e8..081fb520d 100644
--- a/tests/unit_tests/test_plotter.py
+++ b/tests/unit_tests/test_plotter.py
@@ -1,22 +1,27 @@
 import numpy as np
 import openmc
 import pytest
+from matplotlib.figure import Figure
 
 
-def test_calculate_cexs_elem_mat_sab():
-    """Checks that sab cross sections are included in the
-    _calculate_cexs_elem_mat method and have the correct shape"""
-
+@pytest.fixture(scope="module")
+def test_mat():
     mat_1 = openmc.Material()
     mat_1.add_element("H", 4.0, "ao")
     mat_1.add_element("O", 4.0, "ao")
     mat_1.add_element("C", 4.0, "ao")
+    return mat_1
 
-    mat_1.add_s_alpha_beta("c_C6H6")
-    mat_1.set_density("g/cm3", 0.865)
+
+def test_calculate_cexs_elem_mat_sab(test_mat):
+    """Checks that sab cross sections are included in the
+    _calculate_cexs_elem_mat method and have the correct shape"""
+
+    test_mat.add_s_alpha_beta("c_C6H6")
+    test_mat.set_density("g/cm3", 0.865)
 
     energy_grid, data = openmc.plotter._calculate_cexs_elem_mat(
-        mat_1,
+        test_mat,
         ["inelastic"],
         sab_name="c_C6H6",
     )
@@ -28,7 +33,7 @@ def test_calculate_cexs_elem_mat_sab():
     assert len(data[0]) == len(energy_grid)
 
 
-@pytest.mark.parametrize("this, data_type", [("Li", "element"), ("Li6", "nuclide")])
+@pytest.mark.parametrize("this,data_type", [("Li", "element"), ("Li6", "nuclide")])
 def test_calculate_cexs_with_element(this, data_type):
 
     # single type (reaction)
@@ -55,3 +60,28 @@ def test_calculate_cexs_with_element(this, data_type):
     assert len(data[0]) == len(energy_grid)
     # reactions are both the same MT number 2 is elastic
     assert np.array_equal(data[0], data[1])
+
+
+def test_calculate_cexs_with_materials(test_mat):
+    energy_grid, data = openmc.plotter.calculate_cexs(
+        this=test_mat, types=[205], data_type="material"
+    )
+
+    assert isinstance(energy_grid, np.ndarray)
+    assert isinstance(data, np.ndarray)
+    assert len(energy_grid) > 1
+    assert len(data) == 1
+    assert len(data[0]) == len(energy_grid)
+
+
+@pytest.mark.parametrize(("this,data_type"), [("Be", "element"), ("Be9", "nuclide")])
+def test_plot_xs(this, data_type):
+    assert isinstance(
+        openmc.plotter.plot_xs(this, data_type=data_type, types=["total"]), Figure
+    )
+
+
+def test_plot_xs_mat(test_mat):
+    assert isinstance(
+        openmc.plotter.plot_xs(test_mat, data_type="material", types=["total"]), Figure
+    )