Fix #481: make tests of Pubchem more robust

dev_tools/conf/pytest.ini

@@ -15,6 +15,10 @@
 [pytest]
 markers =
     slow: marks tests as slow (deselect with '-m "not slow"')
+    integration: marks tests of integrations with other systems
 
 filterwarnings =
     ignore:Skipped assert_qasm_is_consistent_with_unitary because qiskit.*:UserWarning
+
+# Add this line to always exclude tests marked as 'integration' by default
+addopts = -m "not integration"

src/openfermion/chem/pubchem_test.py

@@ -9,22 +9,84 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+
 """Tests for pubchem.py."""
 
+import time
 import unittest
+from unittest.mock import patch
+
 import numpy
 import pytest
 
 from openfermion.chem.pubchem import geometry_from_pubchem
 from openfermion.testing.testing_utils import module_importable
 
+
+class MockCompound:
+    def __init__(self, atoms):
+        self._atoms = atoms
+
+    def to_dict(self, properties):
+        return {'atoms': self._atoms}
+
+
+def mock_get_compounds(name, searchtype, record_type='2d'):
+    if name == 'water' and record_type == '3d':
+        return [
+            MockCompound(
+                [
+                    {'aid': 1, 'number': 8, 'element': 'O', 'y': 0, 'z': 0, 'x': 0},
+                    {'aid': 2, 'number': 1, 'element': 'H', 'y': 0.8929, 'z': 0.2544, 'x': 0.2774},
+                    {
+                        'aid': 3,
+                        'number': 1,
+                        'element': 'H',
+                        'y': -0.2383,
+                        'z': -0.7169,
+                        'x': 0.6068,
+                    },
+                ]
+            )
+        ]
+    if name == 'water' and record_type == '2d':
+        return [
+            MockCompound(
+                [
+                    {'aid': 1, 'number': 8, 'element': 'O', 'y': -0.155, 'x': 2.5369},
+                    {'aid': 2, 'number': 1, 'element': 'H', 'y': 0.155, 'x': 3.0739},
+                    {'aid': 3, 'number': 1, 'element': 'H', 'y': 0.155, 'x': 2},
+                ]
+            )
+        ]
+    if name == 'helium' and record_type == '2d':
+        return [MockCompound([{'aid': 1, 'number': 2, 'element': 'He', 'y': 0, 'x': 2}])]
+    return []
+
+
 using_pubchempy = pytest.mark.skipif(
     module_importable('pubchempy') is False, reason='Not detecting `pubchempy`.'
 )
 
 
 @using_pubchempy
 class OpenFermionPubChemTest(unittest.TestCase):
+    def _get_geometry_with_retries(self, name):
+        import urllib.error
+
+        import pubchempy
+
+        max_retries = 3
+        delay = 2
+        for attempt in range(max_retries):
+            try:
+                return geometry_from_pubchem(name)
+            except (pubchempy.PubChemHTTPError, urllib.error.URLError):
+                if attempt == max_retries - 1:
+                    raise
+                time.sleep(delay)
+
+    @patch('pubchempy.get_compounds', mock_get_compounds)
     def test_water(self):
         water_geometry = geometry_from_pubchem('water')
         self.water_natoms = len(water_geometry)
@@ -64,18 +126,21 @@ def test_water(self):
         self.assertTrue(water_bond_angle_low <= self.water_bond_angle)
         self.assertTrue(water_bond_angle_high >= self.water_bond_angle)
 
+    @patch('pubchempy.get_compounds', mock_get_compounds)
     def test_helium(self):
         helium_geometry = geometry_from_pubchem('helium')
         self.helium_natoms = len(helium_geometry)
 
         helium_natoms = 1
         self.assertEqual(helium_natoms, self.helium_natoms)
 
+    @patch('pubchempy.get_compounds', mock_get_compounds)
     def test_none(self):
         none_geometry = geometry_from_pubchem('none')
 
         self.assertIsNone(none_geometry)
 
+    @patch('pubchempy.get_compounds', mock_get_compounds)
     def test_water_2d(self):
         water_geometry = geometry_from_pubchem('water', structure='2d')
         self.water_natoms = len(water_geometry)
@@ -91,3 +156,60 @@ def test_water_2d(self):
 
         with pytest.raises(ValueError, match='Incorrect value for the argument structure'):
             _ = geometry_from_pubchem('water', structure='foo')
+
+    @pytest.mark.integration
+    def test_geometry_from_pubchem_live_api(self):
+        water_geometry = self._get_geometry_with_retries('water')  # pragma: no cover
+        self.assertEqual(len(water_geometry), 3)  # pragma: no cover
+
+    @patch('time.sleep', return_value=None)
+    @patch('pubchempy.get_compounds')
+    def test_geometry_from_pubchem_retry_success(self, mock_get_compounds, mock_sleep):
+        import pubchempy
+
+        mock_get_compounds.side_effect = [
+            pubchempy.PubChemHTTPError(503, 'Server Busy', 'Testing'),
+            pubchempy.PubChemHTTPError(503, 'Server Busy', 'Testing'),
+            [
+                MockCompound(
+                    [
+                        {'aid': 1, 'number': 8, 'element': 'O', 'y': 0, 'z': 0, 'x': 0},
+                        {
+                            'aid': 2,
+                            'number': 1,
+                            'element': 'H',
+                            'y': 0.8929,
+                            'z': 0.2544,
+                            'x': 0.2774,
+                        },
+                        {
+                            'aid': 3,
+                            'number': 1,
+                            'element': 'H',
+                            'y': -0.2383,
+                            'z': -0.7169,
+                            'x': 0.6068,
+                        },
+                    ]
+                )
+            ],
+        ]
+
+        water_geometry = self._get_geometry_with_retries('water')
+
+        self.assertEqual(len(water_geometry), 3)
+        self.assertEqual(mock_get_compounds.call_count, 3)
+        self.assertEqual(mock_sleep.call_count, 2)
+
+    @patch('time.sleep', return_value=None)
+    @patch('pubchempy.get_compounds')
+    def test_geometry_from_pubchem_retry_failure(self, mock_get_compounds, mock_sleep):
+        import pubchempy
+
+        mock_get_compounds.side_effect = pubchempy.PubChemHTTPError(503, 'Server Busy', 'Testing')
+
+        with self.assertRaises(pubchempy.PubChemHTTPError):
+            self._get_geometry_with_retries('water')
+
+        self.assertEqual(mock_get_compounds.call_count, 3)
+        self.assertEqual(mock_sleep.call_count, 2)