Added some tests about the electronic structure ground state

crystals/atom.py

@@ -545,6 +545,9 @@ def __init__(self, shells: Dict[Union[str, Orbital], int]):
         for k, v in shells.items():
             self.__setitem__(k, v)
 
+    def __iter__(self):
+        yield from self._structure
+
     def __setitem__(self, key: Union[str, Orbital], value: int):
         # We check that the number of electrons in each Orbital does not
         # go above maximum possible.

crystals/tests/test_atom.py

@@ -222,6 +222,19 @@ def test_orbital_maximum_electrons():
         assert maxima[shell.value[-1]] == Orbital.maximum_electrons(shell)
 
 
+@pytest.mark.parametrize("element, structure", (["C", {"1s": 2, "2s": 2, "2p": 2}],))
+def test_electronic_structure_ground_state_examples(element, structure):
+    gs = ElectronicStructure.ground_state(element)
+    assert gs == ElectronicStructure(shells=structure)
+
+
+@pytest.mark.parametrize("element", range(1, 24))
+def test_electronic_structure_ground_state_num_electrons(element):
+    gs = ElectronicStructure.ground_state(element)
+    num_electrons = sum(gs[k] for k in gs)
+    assert num_electrons == element
+
+
 def test_electronic_structure_maximum_electrons():
     """Test that an error is raised for impossible electronic structures."""
     with pytest.raises(ValueError):