Address review comments, couple of more asserts

mosdef-hub · Apr 28, 2022 · fbbfdd4 · fbbfdd4
1 parent 279ca7e
commit fbbfdd4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/gmso/core/forcefield.py b/gmso/core/forcefield.py
@@ -118,9 +118,8 @@ def non_element_types(self):
             )  # FixMe: Should we make this a first class citizen?
             if element_symbol:
                 element = element_by_symbol(element_symbol)
-                non_element_types.add(None if element else element_symbol)
+                non_element_types.add(element_symbol) if not element else None
 
-        non_element_types.discard(None)
         return non_element_types
 
     @property

diff --git a/gmso/tests/test_forcefield.py b/gmso/tests/test_forcefield.py
@@ -569,3 +569,4 @@ def test_non_element_types(self, non_element_ff, opls_ethane_foyer):
         assert "_CH3" in non_element_ff.non_element_types
         assert "_CH2" in non_element_ff.non_element_types
         assert opls_ethane_foyer.non_element_types == set()
+        assert len(opls_ethane_foyer.atom_types) > 0