Skip to content

Commit c94825f

Browse files
committed
Fix #481: add mocking of Pubchem network calls
In order to avoid failures in CI, this follows a [recommendation from Pavol Juhas on PR #1113]( #1113 (review)) about using mocking to avoid flaky test behavior in `pubchem_test.py`.
1 parent fd6456e commit c94825f

1 file changed

Lines changed: 149 additions & 0 deletions

File tree

src/openfermion/chem/pubchem_test.py

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,20 +11,80 @@
1111
# limitations under the License.
1212
"""Tests for pubchem.py."""
1313

14+
import time
1415
import unittest
16+
from unittest.mock import patch
17+
1518
import numpy
1619
import pytest
1720

1821
from openfermion.chem.pubchem import geometry_from_pubchem
1922
from openfermion.testing.testing_utils import module_importable
2023

24+
25+
class MockCompound:
26+
def __init__(self, atoms):
27+
self._atoms = atoms
28+
29+
def to_dict(self, properties):
30+
return {'atoms': self._atoms}
31+
32+
33+
def mock_get_compounds(name, searchtype, record_type='2d'):
34+
if name == 'water' and record_type == '3d':
35+
return [
36+
MockCompound(
37+
[
38+
{'aid': 1, 'number': 8, 'element': 'O', 'y': 0, 'z': 0, 'x': 0},
39+
{'aid': 2, 'number': 1, 'element': 'H', 'y': 0.8929, 'z': 0.2544, 'x': 0.2774},
40+
{
41+
'aid': 3,
42+
'number': 1,
43+
'element': 'H',
44+
'y': -0.2383,
45+
'z': -0.7169,
46+
'x': 0.6068,
47+
},
48+
]
49+
)
50+
]
51+
if name == 'water' and record_type == '2d':
52+
return [
53+
MockCompound(
54+
[
55+
{'aid': 1, 'number': 8, 'element': 'O', 'y': -0.155, 'x': 2.5369},
56+
{'aid': 2, 'number': 1, 'element': 'H', 'y': 0.155, 'x': 3.0739},
57+
{'aid': 3, 'number': 1, 'element': 'H', 'y': 0.155, 'x': 2},
58+
]
59+
)
60+
]
61+
if name == 'helium' and record_type == '2d':
62+
return [MockCompound([{'aid': 1, 'number': 2, 'element': 'He', 'y': 0, 'x': 2}])]
63+
return []
64+
65+
2166
using_pubchempy = pytest.mark.skipif(
2267
module_importable('pubchempy') is False, reason='Not detecting `pubchempy`.'
2368
)
2469

2570

2671
@using_pubchempy
2772
class OpenFermionPubChemTest(unittest.TestCase):
73+
def _get_geometry_with_retries(self, name):
74+
import pubchempy
75+
import urllib.error
76+
77+
max_retries = 3
78+
delay = 2
79+
for attempt in range(max_retries):
80+
try:
81+
return geometry_from_pubchem(name)
82+
except (pubchempy.PubChemHTTPError, urllib.error.URLError):
83+
if attempt == max_retries - 1:
84+
raise
85+
time.sleep(delay)
86+
87+
@patch('pubchempy.get_compounds', mock_get_compounds)
2888
def test_water(self):
2989
water_geometry = geometry_from_pubchem('water')
3090
self.water_natoms = len(water_geometry)
@@ -64,18 +124,21 @@ def test_water(self):
64124
self.assertTrue(water_bond_angle_low <= self.water_bond_angle)
65125
self.assertTrue(water_bond_angle_high >= self.water_bond_angle)
66126

127+
@patch('pubchempy.get_compounds', mock_get_compounds)
67128
def test_helium(self):
68129
helium_geometry = geometry_from_pubchem('helium')
69130
self.helium_natoms = len(helium_geometry)
70131

71132
helium_natoms = 1
72133
self.assertEqual(helium_natoms, self.helium_natoms)
73134

135+
@patch('pubchempy.get_compounds', mock_get_compounds)
74136
def test_none(self):
75137
none_geometry = geometry_from_pubchem('none')
76138

77139
self.assertIsNone(none_geometry)
78140

141+
@patch('pubchempy.get_compounds', mock_get_compounds)
79142
def test_water_2d(self):
80143
water_geometry = geometry_from_pubchem('water', structure='2d')
81144
self.water_natoms = len(water_geometry)
@@ -91,3 +154,89 @@ def test_water_2d(self):
91154

92155
with pytest.raises(ValueError, match='Incorrect value for the argument structure'):
93156
_ = geometry_from_pubchem('water', structure='foo')
157+
158+
def test_geometry_from_pubchem_live_api(self):
159+
try:
160+
import pubchempy
161+
except ImportError: # pragma: no cover
162+
return
163+
164+
water_geometry = self._get_geometry_with_retries('water')
165+
self.assertEqual(len(water_geometry), 3)
166+
167+
@patch('time.sleep', return_value=None)
168+
@patch('pubchempy.get_compounds')
169+
def test_geometry_from_pubchem_retry_success(self, mock_get_compounds, mock_sleep):
170+
try:
171+
import pubchempy
172+
except ImportError: # pragma: no cover
173+
return
174+
175+
# Simulate HTTP error with proper code attribute for the first 2 calls, then succeed
176+
class FakeHTTPError(Exception):
177+
def __init__(self, code, msg):
178+
self.code = code
179+
self.reason = msg
180+
181+
def read(self):
182+
return b'{"Fault": {"Details": ["busy"]}}'
183+
184+
mock_get_compounds.side_effect = [
185+
pubchempy.PubChemHTTPError(FakeHTTPError(503, 'Server Busy')),
186+
pubchempy.PubChemHTTPError(FakeHTTPError(503, 'Server Busy')),
187+
[
188+
MockCompound(
189+
[
190+
{'aid': 1, 'number': 8, 'element': 'O', 'y': 0, 'z': 0, 'x': 0},
191+
{
192+
'aid': 2,
193+
'number': 1,
194+
'element': 'H',
195+
'y': 0.8929,
196+
'z': 0.2544,
197+
'x': 0.2774,
198+
},
199+
{
200+
'aid': 3,
201+
'number': 1,
202+
'element': 'H',
203+
'y': -0.2383,
204+
'z': -0.7169,
205+
'x': 0.6068,
206+
},
207+
]
208+
)
209+
],
210+
]
211+
212+
water_geometry = self._get_geometry_with_retries('water')
213+
214+
self.assertEqual(len(water_geometry), 3)
215+
self.assertEqual(mock_get_compounds.call_count, 3)
216+
self.assertEqual(mock_sleep.call_count, 2)
217+
218+
@patch('time.sleep', return_value=None)
219+
@patch('pubchempy.get_compounds')
220+
def test_geometry_from_pubchem_retry_failure(self, mock_get_compounds, mock_sleep):
221+
try:
222+
import pubchempy
223+
except ImportError: # pragma: no cover
224+
return
225+
226+
class FakeHTTPError(Exception):
227+
def __init__(self, code, msg):
228+
self.code = code
229+
self.reason = msg
230+
231+
def read(self):
232+
return b'{"Fault": {"Details": ["busy"]}}'
233+
234+
mock_get_compounds.side_effect = pubchempy.PubChemHTTPError(
235+
FakeHTTPError(503, 'Server Busy')
236+
)
237+
238+
with self.assertRaises(pubchempy.PubChemHTTPError):
239+
self._get_geometry_with_retries('water')
240+
241+
self.assertEqual(mock_get_compounds.call_count, 3)
242+
self.assertEqual(mock_sleep.call_count, 2)

0 commit comments

Comments
 (0)