add test for site level

mosdef-hub · Nov 13, 2023 · 2c17ff9 · 2c17ff9
1 parent b243619
commit 2c17ff9
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/gmso/tests/test_atom.py b/gmso/tests/test_atom.py
@@ -74,6 +74,12 @@ def test_bad_pos_input(self, position):
         with pytest.raises((u.exceptions.InvalidUnitOperation, ValueError)):
             Atom(name="atom", position=u.nm * position)
 
+    def test_pos_restraint(self, spce_water):
+        for site in spce_water.sites:
+            for ax in ["kx", "ky", "kz"]:
+                assert site.restraint[ax] == 1000 * u.Unit("kJ/(mol*nm**2)")
+                site.restraint[ax] = 500 * u.Unit("kJ/(mol*nm**2)")
+
     def test_equivalence(self):
         ref = Atom(name="atom", position=u.nm * np.zeros(3))
         same_atom = Atom(name="atom", position=u.nm * np.zeros(3))