add a test case for the argument 'structure' in geometry_from_pubchem

quantumlib · Nov 28, 2019 · c7f2b05 · c7f2b05
1 parent 7cbd575
commit c7f2b05
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/src/openfermion/utils/_pubchem_test.py b/src/openfermion/utils/_pubchem_test.py
@@ -76,3 +76,16 @@ def test_none(self):
         none_geometry = geometry_from_pubchem('none')
 
         self.assertIsNone(none_geometry)
+
+    def test_water_2d(self):
+        water_geometry = geometry_from_pubchem('water', structure='2d')
+        self.water_natoms = len(water_geometry)
+
+        water_natoms = 3
+        self.assertEqual(water_natoms, self.water_natoms)
+
+        self.oxygen_z_1 = water_geometry[0][1][2]
+        self.oxygen_z_2 = water_geometry[1][1][2]
+        z = 0
+        self.assertEqual(z, self.oxygen_z_1)
+        self.assertEqual(z, self.oxygen_z_2)