async def test_files_to_retrieve_cut_short(mocker, mock_sleep):
def make_header_mocks(init):
for i in range(100):
# first time needs headers for lastmod
if i == 0:
hmock = MagicMock()
hmock.headers.__getitem__.return_value = init
yield hmock
elif i == 10:
yield aiohttp.ClientOSError
else:
yield i
model = nwp.RAP.copy()
# so now is after avg max run length
model['avg_max_run_length'] = '30min'
init = pd.Timestamp('20190409T0000Z')
headers = list(make_header_mocks(init))
session = MagicMock()
session.head.return_value.__aenter__.side_effect = headers
checknext = mocker.patch(
'solarforecastarbiter.io.fetch.nwp.check_next_inittime',
new=CoroutineMock())
checknext.return_value = True
params = [p async for p in nwp.files_to_retrieve(session, model, Path('/'),
init)]
assert len(params) == 10
async def test_files_to_retrieve(mocker, mock_sleep):
session = MagicMock()
session.head.return_value.__aenter__.return_value.headers.__getitem__.return_value = '20190101T0000Z' # NOQA
params = [p async for p in nwp.files_to_retrieve(
session, nwp.RAP, Path('/'), pd.Timestamp('20190409T1200Z'))]
assert len(params) == 22