From 98253ab9b4e2bc3a8cf7805ddc049adee235c17d Mon Sep 17 00:00:00 2001
From: Mikolaj Adam Kowalski <mak60@cam.ac.uk>
Date: Fri, 14 Feb 2020 17:54:01 +0000
Subject: [PATCH] Implemented unit tests for Element.expand

---
 openmc/element.py                | 20 ++++++---
 tests/unit_tests/test_element.py | 75 ++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 6 deletions(-)
 create mode 100644 tests/unit_tests/test_element.py

diff --git a/openmc/element.py b/openmc/element.py
index 4bedda77d..dd3ef1535 100644
--- a/openmc/element.py
+++ b/openmc/element.py
@@ -54,8 +54,8 @@ class Element(str):
             'ao' for atom percent and 'wo' for weight percent
         enrichment : float, optional
             Enrichment of an enrichment_taget nuclide in percent (ao or wo).
-            If enrichment_taget is not supplied then it is enrichment for U235 in
-            weight percent. For example, input 4.95 for 4.95 weight percent
+            If enrichment_taget is not supplied then it is enrichment for U235
+            in weight percent. For example, input 4.95 for 4.95 weight percent
             enriched U. Default is None (natural composition).
         enrichment_target: str, optional
             Single nuclide name to enrich from a natural composition e.g. O16
@@ -86,8 +86,8 @@ class Element(str):
             of the element is requested
 
         ValueError
-            Enrichment is requested of the element composed of more or less
-            then 2 isotopes.
+            Enrichment is requested of the element that is not composed of
+            two isotopes.
 
         Notes
         -----
@@ -101,6 +101,8 @@ class Element(str):
         Function does not check if enrichment value is in a valid range <0;100>
 
         """
+        # Check input
+        cv.check_value('enrichment_type', enrichment_type, {'ao', 'wo'})
 
         # Get the nuclides present in nature
         natural_nuclides = set()
@@ -193,6 +195,12 @@ class Element(str):
         # Old treatment for Uranium
         if enrichment is not None and enrichment_target is None:
 
+            # Check that the element is Uranium
+            if self.name != 'U':
+                msg = 'Enrichment procedure for Uranium was requested, '\
+                      'but the isotope is {0} not U'.format(self)
+                raise ValueError(msg)
+
             # Calculate the mass fractions of isotopes
             abundances['U234'] = 0.0089 * enrichment
             abundances['U235'] = enrichment
@@ -209,8 +217,8 @@ class Element(str):
                 abundances[nuclide] /= sum_abundances
 
         # Modify mole fractions if enrichment provided
-        # New treatment for arbitrary element94.96165869352852
-        # Interpret required enrichment as weight %
+        # New treatment for arbitrary element
+        # Interpret required enrichment as weight
         elif enrichment is not None and enrichment_target is not None:
 
             # Check if is a single isotope mixture
diff --git a/tests/unit_tests/test_element.py b/tests/unit_tests/test_element.py
new file mode 100644
index 000000000..1143b5357
--- /dev/null
+++ b/tests/unit_tests/test_element.py
@@ -0,0 +1,75 @@
+import openmc.element
+import pytest as pt
+
+from openmc.data import NATURAL_ABUNDANCE, atomic_mass
+
+# Relative tolerance for float comparison
+TOL = 1e-9
+
+
+def test_expand_no_ernichment():
+    """ Expand Li in natural compositions"""
+    lithium = openmc.element.Element('Li')
+
+    # Verify the expansion into ATOMIC fraction against natural composition
+    for isotope in lithium.expand(100.0, 'ao'):
+        assert isotope[1] == pt.approx(NATURAL_ABUNDANCE[isotope[0]] * 100.0, rel=TOL)
+
+    # Verify the expanction into WEIGHT fraction against natural composition
+    natural = {'Li6': NATURAL_ABUNDANCE['Li6'] * atomic_mass('Li6'),
+               'Li7': NATURAL_ABUNDANCE['Li7'] * atomic_mass('Li7')}
+    li_am = sum(natural.values())
+    for key in natural.keys():
+        natural[key] /= li_am
+
+    for isotope in lithium.expand(100.0, 'wo'):
+        assert isotope[1] == pt.approx(natural[isotope[0]] * 100.0, rel=TOL)
+
+
+def test_expand_enrichment():
+    """ Expand and verify enrichment of Li """
+    lithium = openmc.element.Element('Li')
+
+    # Verify the enrichment by atoms
+    ref = {'Li6': 75.0, 'Li7': 25.0}
+    for isotope in lithium.expand(100.0, 'ao', 25.0, 'Li7', 'ao'):
+        assert isotope[1] == pt.approx(ref[isotope[0]], rel=TOL)
+
+    # Verify the enrichment by weight
+    for isotope in lithium.expand(100.0, 'wo', 25.0, 'Li7', 'wo'):
+        assert isotope[1] == pt.approx(ref[isotope[0]], rel=TOL)
+
+
+def test_expand_exceptions():
+    """ Test that correct exceptions are raiused for invalid input """
+
+    # 1 Isotope Element
+    with pt.raises(ValueError):
+        element = openmc.element.Element('Be')
+        fail = element.expand(70.0, 'ao', 4.0, 'Be9')
+
+    # 3 Isotope Element
+    with pt.raises(ValueError):
+        element = openmc.element.Element('Cr')
+        fail = element.expand(70.0, 'ao', 4.0, 'Cr52')
+
+    # Non-present Enrichment Target
+    with pt.raises(ValueError):
+        element = openmc.element.Element('H')
+        fail = element.expand(70.0, 'ao', 4.0, 'H4')
+
+    # Enrichment Procedure for Uranium if not Uranium
+    with pt.raises(ValueError):
+        element = openmc.element.Element('Li')
+        fail = element.expand(70.0, 'ao', 4.0)
+
+    # Missing Enrichment Target
+    with pt.raises(ValueError):
+        element = openmc.element.Element('Li')
+        fail = element.expand(70.0, 'ao', 4.0, enrichment_type='ao')
+
+    # Invalid Enrichment Type Entry
+    with pt.raises(ValueError):
+        element = openmc.element.Element('Li')
+        fail = element.expand(70.0, 'ao', 4.0, 'Li7','Grand Moff Tarkin')
+