1111# limitations under the License.
1212"""Tests for pubchem.py."""
1313
14+ import time
1415import unittest
16+ from unittest .mock import patch
17+
1518import numpy
1619import pytest
1720
1821from openfermion .chem .pubchem import geometry_from_pubchem
1922from 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+
2166using_pubchempy = pytest .mark .skipif (
2267 module_importable ('pubchempy' ) is False , reason = 'Not detecting `pubchempy`.'
2368)
2469
2570
2671@using_pubchempy
2772class 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