Commit 91d5e0eb authored by Derek Homeier's avatar Derek Homeier
Browse files

Fetch with default filenames, expand obsrec into rows

parent 359f7162
Pipeline #2745 failed with stage
in 3 minutes and 45 seconds
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
import warnings
import os import os
import copy
import warnings
import astropy.units as u import astropy.units as u
from astropy.time import Time from astropy.time import Time
...@@ -274,7 +275,13 @@ class KISClient(BaseClient): ...@@ -274,7 +275,13 @@ class KISClient(BaseClient):
try: try:
response = urllib.request.urlopen(full_url) response = urllib.request.urlopen(full_url)
obs_dict = json.loads(response.read()).get('_embedded', []) obs_dict = json.loads(response.read()).get('_embedded', [])
results += obs_dict if len(obs_dict) > 0 and 'links' in obs_dict[0]:
l1_data = obs_dict[0]['links'].get('l1_data', [])
obs_rec = []
for exp in l1_data:
obs_rec.append(copy.deepcopy(obs_dict[0]))
obs_rec[-1]['links']['l1_data'] = exp
results += obs_rec
except(HTTPError, URLError) as exc: except(HTTPError, URLError) as exc:
raise URLError(f'Unable to execute search "{full_url}": {exc}. Confirm that ' raise URLError(f'Unable to execute search "{full_url}": {exc}. Confirm that '
f'RESTHeart is running on {self._BASE_URL} and connected.') f'RESTHeart is running on {self._BASE_URL} and connected.')
...@@ -310,12 +317,13 @@ class KISClient(BaseClient): ...@@ -310,12 +317,13 @@ class KISClient(BaseClient):
for row in query_results: for row in query_results:
inst = row['description']['INSTRUMENT'] inst = row['description']['INSTRUMENT']
rowpath = os.path.join(path, row['_id']['$oid']) oid = row['links']['l1_data']['$oid']
for l1_data in row['links']['l1_data']: # Content-Disposition header default is "{row['_id']['$oid']}/{oid}.{ext}" (no '.json').
oid = l1_data['$oid'] # rowpath = row['_id']['$oid']
filename = f"{oid}.{ext}" filepath = os.path.join(row['description']['OBS_NAME'], f"{oid}.{ext}")
url = f"{self._BASE_URL}{inst}_l1_data.files/{oid}{binfile}" url = f"{self._BASE_URL}{inst}_l1_data.files/{oid}{binfile}"
downloader.enqueue_file(url, filename=os.path.join(rowpath, filename), max_splits=1) downloader.enqueue_file(url, filename=str(path).format(file=filepath,
**row.response_block_map))
@classmethod @classmethod
def _can_handle_query(cls, *query, hasinstr=False): def _can_handle_query(cls, *query, hasinstr=False):
......
...@@ -2,6 +2,7 @@ import pytest ...@@ -2,6 +2,7 @@ import pytest
import urllib.request import urllib.request
from urllib.error import HTTPError, URLError from urllib.error import HTTPError, URLError
from pathlib import Path
import os import os
import json import json
import parfive import parfive
...@@ -9,6 +10,7 @@ import parfive ...@@ -9,6 +10,7 @@ import parfive
import astropy.units as u import astropy.units as u
from astropy.io import fits from astropy.io import fits
from astropy.time import Time from astropy.time import Time
import sunpy
from sunpy.net import Fido from sunpy.net import Fido
from sunpy.net import attrs as a from sunpy.net import attrs as a
from sunpy.net.base_client import QueryResponseTable from sunpy.net.base_client import QueryResponseTable
...@@ -31,6 +33,10 @@ except(HTTPError, URLError): ...@@ -31,6 +33,10 @@ except(HTTPError, URLError):
HAS_DOCKERTEST = False HAS_DOCKERTEST = False
def dirnames(path):
return os.path.dirname(path).split(os.path.sep)
@pytest.fixture @pytest.fixture
def client(): def client():
return KISClient() return KISClient()
...@@ -62,7 +68,7 @@ def test_docker(client): ...@@ -62,7 +68,7 @@ def test_docker(client):
res = client.search(a.Instrument("GRIS") & a.sdc.ObsName('gris_20140426_000')) res = client.search(a.Instrument("GRIS") & a.sdc.ObsName('gris_20140426_000'))
assert isinstance(res, QueryResponseTable) assert isinstance(res, QueryResponseTable)
assert len(res) == 1 assert len(res) == 105
description = res[0].get('description') description = res[0].get('description')
assert len(description) == 34 assert len(description) == 34
assert description['INSTRUMENT'] == 'gris' assert description['INSTRUMENT'] == 'gris'
...@@ -71,7 +77,7 @@ def test_docker(client): ...@@ -71,7 +77,7 @@ def test_docker(client):
assert description['DATE_BEG']['$date'] == 1398505619000 assert description['DATE_BEG']['$date'] == 1398505619000
assert description['DATE_END']['$date'] == 1398506021300 assert description['DATE_END']['$date'] == 1398506021300
file_ids = [ld['$oid'] for ld in res[0]['links']['l1_data']] file_ids = [ld['links']['l1_data']['$oid'] for ld in res]
assert len(file_ids) == 105 assert len(file_ids) == 105
for oid in file_ids[0], file_ids[104]: for oid in file_ids[0], file_ids[104]:
meta = json.loads(urllib.request.urlopen(f"{_BASE_URL}gris_l1_data.files/{oid}").read()) meta = json.loads(urllib.request.urlopen(f"{_BASE_URL}gris_l1_data.files/{oid}").read())
...@@ -87,8 +93,8 @@ def test_docker(client): ...@@ -87,8 +93,8 @@ def test_docker(client):
rowpath = f"{res[0]['_id']['$oid']}" rowpath = f"{res[0]['_id']['$oid']}"
binfile = '' binfile = ''
ext = 'json' ext = 'json'
for i, l1_data in enumerate(res[0]['links']['l1_data'][:10]): for i, ld in enumerate(res[:10]):
oid = l1_data['$oid'] oid = ld['links']['l1_data']['$oid']
filename = f"{oid}.{ext}" filename = f"{oid}.{ext}"
url = f"{_BASE_URL}{inst}_l1_data.files/{oid}{binfile}" url = f"{_BASE_URL}{inst}_l1_data.files/{oid}{binfile}"
assert url == f"{_BASE_URL}gris_l1_data.files/{file_ids[i]}" assert url == f"{_BASE_URL}gris_l1_data.files/{file_ids[i]}"
...@@ -96,15 +102,15 @@ def test_docker(client): ...@@ -96,15 +102,15 @@ def test_docker(client):
binfile = '/binary' binfile = '/binary'
ext = 'fits' ext = 'fits'
for l1_data in res[0]['links']['l1_data'][:2]: for ld in res[:2]:
oid = l1_data['$oid'] oid = ld['links']['l1_data']['$oid']
filename = f"{oid}.{ext}" filename = f"{oid}.{ext}"
url = f"{_BASE_URL}{inst}_l1_data.files/{oid}{binfile}" url = f"{_BASE_URL}{inst}_l1_data.files/{oid}{binfile}"
downloader.enqueue_file(url, filename=os.path.join(rowpath, filename), max_splits=1) downloader.enqueue_file(url, filename=os.path.join(rowpath, filename), max_splits=1)
assert downloader.queued_downloads == 12 assert downloader.queued_downloads == 12
assert downloader.http_queue[0].keywords['url'].startswith(_BASE_URL) assert downloader.http_queue[0].keywords['url'].startswith(_BASE_URL)
assert res[0]['links']['l1_data'][0]['$oid'] in downloader.http_queue[0].keywords['url'] assert res[0]['links']['l1_data']['$oid'] in downloader.http_queue[0].keywords['url']
assert downloader.http_queue[10].keywords['url'].endswith(binfile) assert downloader.http_queue[10].keywords['url'].endswith(binfile)
files = downloader.download() files = downloader.download()
...@@ -159,7 +165,7 @@ def test_search(client): ...@@ -159,7 +165,7 @@ def test_search(client):
if HAS_DOCKERTEST: if HAS_DOCKERTEST:
res = client.search(query) res = client.search(query)
assert isinstance(res, QueryResponseTable) assert isinstance(res, QueryResponseTable)
assert len(res) > 50 assert len(res) == 1
assert 'description' in res.colnames assert 'description' in res.colnames
else: else:
with pytest.raises(URLError, match=rf"{_dockerexc('bbi')}" with pytest.raises(URLError, match=rf"{_dockerexc('bbi')}"
...@@ -173,7 +179,7 @@ def test_search(client): ...@@ -173,7 +179,7 @@ def test_search(client):
assert client._can_handle_query(query) assert client._can_handle_query(query)
if HAS_DOCKERTEST: if HAS_DOCKERTEST:
res = client.search(query) res = client.search(query)
assert len(res) == 100 assert len(res) == 1
assert 'description' in res.colnames assert 'description' in res.colnames
else: else:
with pytest.raises(URLError, match=rf"{_dockerexc('lars')}" with pytest.raises(URLError, match=rf"{_dockerexc('lars')}"
...@@ -185,7 +191,7 @@ def test_search(client): ...@@ -185,7 +191,7 @@ def test_search(client):
assert client._can_handle_query(*query) assert client._can_handle_query(*query)
if HAS_DOCKERTEST: if HAS_DOCKERTEST:
res = client.search(*query) res = client.search(*query)
assert len(res) == 100 assert len(res) == 1
assert 'description' in res.colnames assert 'description' in res.colnames
else: else:
with pytest.raises(URLError, match=rf"{_dockerexc('lars')}" with pytest.raises(URLError, match=rf"{_dockerexc('lars')}"
...@@ -197,7 +203,7 @@ def test_search(client): ...@@ -197,7 +203,7 @@ def test_search(client):
assert client._can_handle_query(query) assert client._can_handle_query(query)
if HAS_DOCKERTEST: if HAS_DOCKERTEST:
res = client.search(query) res = client.search(query)
assert len(res) == 200 assert len(res) == 362
assert 'description' in res.colnames assert 'description' in res.colnames
assert 'THETA' in res[0]['description'] assert 'THETA' in res[0]['description']
theta = [obs['description']['THETA'] for obs in res] theta = [obs['description']['THETA'] for obs in res]
...@@ -213,7 +219,7 @@ def test_search(client): ...@@ -213,7 +219,7 @@ def test_search(client):
query = a.Instrument("LARS") | a.Instrument("GRIS"), a.sdc.Theta(85*u.deg, 3000*u.arcmin) query = a.Instrument("LARS") | a.Instrument("GRIS"), a.sdc.Theta(85*u.deg, 3000*u.arcmin)
if HAS_DOCKERTEST: if HAS_DOCKERTEST:
res = client.search(*query) res = client.search(*query)
assert len(res) == 200 assert len(res) == 163
assert 'description' in res.colnames assert 'description' in res.colnames
assert 'THETA' in res[0]['description'] assert 'THETA' in res[0]['description']
theta = [obs['description']['THETA'] for obs in res] theta = [obs['description']['THETA'] for obs in res]
...@@ -231,12 +237,12 @@ def test_fido_search(): ...@@ -231,12 +237,12 @@ def test_fido_search():
two_inst = (a.Instrument("LARS") | a.Instrument("GRIS")) two_inst = (a.Instrument("LARS") | a.Instrument("GRIS"))
if HAS_DOCKERTEST: if HAS_DOCKERTEST:
res = Fido.search(a.Instrument("GRIS") & a.sdc.Theta(50*u.deg, 80*u.deg)) res = Fido.search(a.Instrument("GRIS") & a.sdc.Theta(50*u.deg, 80*u.deg))
assert len(res['kis']) == 100 assert len(res['kis']) == 162
theta = [obs['description']['THETA'] for obs in res['kis']] theta = [obs['description']['THETA'] for obs in res['kis']]
assert (min(theta) >= 50) & (max(theta) <= 80) assert (min(theta) >= 50) & (max(theta) <= 80)
res = Fido.search(a.Instrument("GRIS"), a.sdc.Theta(50*u.deg, 80*u.deg)) res = Fido.search(a.Instrument("GRIS"), a.sdc.Theta(50*u.deg, 80*u.deg))
assert len(res['kis']) == 100 assert len(res['kis']) == 162
theta = [obs['description']['THETA'] for obs in res['kis']] theta = [obs['description']['THETA'] for obs in res['kis']]
assert (min(theta) >= 50) & (max(theta) <= 80) assert (min(theta) >= 50) & (max(theta) <= 80)
...@@ -250,7 +256,8 @@ def test_fido_search(): ...@@ -250,7 +256,8 @@ def test_fido_search():
res = Fido.search(two_inst, a.sdc.Theta(50*u.deg, 80*u.deg)) res = Fido.search(two_inst, a.sdc.Theta(50*u.deg, 80*u.deg))
assert len(res['kis']) == 2 assert len(res['kis']) == 2
assert len(res['kis'][0]) == 100 assert len(res['kis'][0]) == 1
assert len(res['kis'][1]) == 162
assert res['kis'][0][0]['description']['INSTRUMENT'] == 'lars' assert res['kis'][0][0]['description']['INSTRUMENT'] == 'lars'
theta = [obs['description']['THETA'] for obs in res['kis'][0]] theta = [obs['description']['THETA'] for obs in res['kis'][0]]
assert (min(theta) >= 50) & (max(theta) <= 80) assert (min(theta) >= 50) & (max(theta) <= 80)
...@@ -260,7 +267,7 @@ def test_fido_search(): ...@@ -260,7 +267,7 @@ def test_fido_search():
date = a.Time("2016/08/26 16:25", "2016/08/26 16:45") date = a.Time("2016/08/26 16:25", "2016/08/26 16:45")
res = Fido.search(a.Instrument("GRIS"), a.sdc.PolStates('iquv'), date) res = Fido.search(a.Instrument("GRIS"), a.sdc.PolStates('iquv'), date)
assert len(res['kis']) == 1 assert len(res['kis']) == 400
assert res['kis'][0]['description']['POL_STATES'] == 'IQUV' assert res['kis'][0]['description']['POL_STATES'] == 'IQUV'
else: else:
with pytest.raises(URLError, match=rf"{_dockerexc('gris')}" with pytest.raises(URLError, match=rf"{_dockerexc('gris')}"
...@@ -284,20 +291,26 @@ def test_fido_fetch(): ...@@ -284,20 +291,26 @@ def test_fido_fetch():
date = a.Time("2017/05/22 08:45", "2017/05/22 08:55") date = a.Time("2017/05/22 08:45", "2017/05/22 08:55")
res = Fido.search(a.Instrument("BBI"), date) res = Fido.search(a.Instrument("BBI"), date)
assert len(res['kis']) == 8 assert len(res['kis']) == 10
desc = [r['description'] for r in res['kis']]
assert max([d['DATE_BEG']['$date'] for d in desc]) < date.end.unix * 1000
assert min([d['DATE_END']['$date'] for d in desc]) > date.start.unix * 1000
files = Fido.fetch(res['kis']) files = Fido.fetch(res['kis'])
assert len(files) == 45 assert len(files) == 10
for filepath in files: for i, filepath in enumerate(files):
assert dirnames(filepath)[-1] == desc[::-1][i]['OBS_NAME']
meta = json.load(open(filepath)) meta = json.load(open(filepath))
assert meta['_id']['$oid'] == os.path.splitext(os.path.basename(filepath))[0] assert meta['_id']['$oid'] == os.path.splitext(os.path.basename(filepath))[0]
assert date.start.isot[:12] in meta['metadata']['header']['DATE-BEG'] assert date.start.isot[:12] in meta['metadata']['header']['DATE-BEG']
assert date.start < Time(meta['metadata']['header']['DATE-BEG']) < date.end assert date.start < Time(meta['metadata']['header']['DATE-BEG']) < date.end
assert desc[i]['OBS_NAME'] in meta['metadata']['header']['FILENAME']
files = Fido.fetch(res['kis'][:1], binary=True) files = Fido.fetch(res['kis'][:3], binary=True)
assert os.path.dirname(files[0]).split(os.path.sep)[-1] == res['kis'][0]['_id']['$oid'] assert len(files) == 3
assert len(files) == 10 for i, filepath in enumerate(files):
for filepath in files: assert dirnames(filepath)[-1] == desc[::-1][i]['OBS_NAME']
hdulist = fits.open(filepath) hdulist = fits.open(filepath)
assert hdulist[0].header.get('TELESCOP') == 'GREGOR' assert hdulist[0].header.get('TELESCOP') == 'GREGOR'
assert hdulist[0].header.get('INSTRUME') == 'BBI' assert hdulist[0].header.get('INSTRUME') == 'BBI'
...@@ -320,13 +333,13 @@ def test_fido_fetch_2(): ...@@ -320,13 +333,13 @@ def test_fido_fetch_2():
else: else:
res = Fido.search((a.Instrument("GRIS") | a.Instrument("LARS")) & date) res = Fido.search((a.Instrument("GRIS") | a.Instrument("LARS")) & date)
assert len(res['kis']) == 2 assert len(res['kis']) == 2
assert len(res['kis'][0]) == 1 assert len(res['kis'][0]) == 400
assert res['kis'][0][0]['description']['DATE_BEG']['$date'] < date.end.unix * 1000 assert res['kis'][0][0]['description']['DATE_BEG']['$date'] < date.end.unix * 1000
assert res['kis'][0][0]['description']['DATE_END']['$date'] > date.start.unix * 1000 assert res['kis'][0][0]['description']['DATE_END']['$date'] > date.start.unix * 1000
files = Fido.fetch(res['kis'], binary=False) files = Fido.fetch(res['kis'][0, :100], binary=False)
assert len(files) == 400 assert len(files) == 100
assert files[0].endswith('.json') assert files[0].endswith('.json')
for filepath in files: for filepath in files:
meta = json.load(open(filepath)) meta = json.load(open(filepath))
...@@ -334,9 +347,21 @@ def test_fido_fetch_2(): ...@@ -334,9 +347,21 @@ def test_fido_fetch_2():
assert date.start.isot[:12] in meta['metadata']['header']['DATE-OBS'] assert date.start.isot[:12] in meta['metadata']['header']['DATE-OBS']
assert date.start < Time(meta['metadata']['header']['DATE-OBS']) assert date.start < Time(meta['metadata']['header']['DATE-OBS'])
files = Fido.fetch(res['kis'][0, :5], binary=True)
assert len(files) == 5
assert files[0].endswith('.fits')
for filepath in files:
hdulist = fits.open(filepath)
assert hdulist[0].header['TELESCOP'] in ('GREGOR', 'VTT')
assert hdulist[0].header['INSTRUME'] in ('GRIS', 'LARS')
assert date.start.iso[:10] in hdulist[0].header['DATE-OBS']
assert Time(hdulist[0].header['DATE-OBS']).mjd < date.end.mjd + 600
assert Time(hdulist[0].header['DATE-OBS']).mjd > date.start.mjd - 600
hdulist.close()
date = a.Time("2016/05/13 10:55", "2016/05/13 11:00") date = a.Time("2016/05/13 10:55", "2016/05/13 11:00")
if not HAS_DOCKERTEST: if not HAS_DOCKERTEST:
with pytest.raises(URLError, match=rf"{_dockerexc('bbi')}" with pytest.raises(URLError, match=rf"{_dockerexc('gris')}"
rf"{{'description.DATE_BEG':{{'.lte':{{'.date':'{date.end.isot}'}}}}}}," rf"{{'description.DATE_BEG':{{'.lte':{{'.date':'{date.end.isot}'}}}}}},"
rf"{{'description.DATE_END':{{'.gte':{{'.date':'{date.start.isot}'}}"): rf"{{'description.DATE_END':{{'.gte':{{'.date':'{date.start.isot}'}}"):
res = Fido.search((a.Instrument("GRIS") | a.Instrument("LARS")) & date) res = Fido.search((a.Instrument("GRIS") | a.Instrument("LARS")) & date)
...@@ -344,39 +369,39 @@ def test_fido_fetch_2(): ...@@ -344,39 +369,39 @@ def test_fido_fetch_2():
res = Fido.search((a.Instrument("GRIS") | a.Instrument("LARS")) & date) res = Fido.search((a.Instrument("GRIS") | a.Instrument("LARS")) & date)
assert len(res['kis']) == 2 assert len(res['kis']) == 2
assert len(res['kis'][0]) == 300
assert len(res['kis'][1]) == 1 assert len(res['kis'][1]) == 1
desc = [r['description'] for r in res['kis'][0]] + [r['description'] for r in res['kis'][1]]
dirs = [d['OBS_NAME'] for d in desc]
assert len(desc) == 301
date_beg = [obs['description']['DATE_BEG']['$date'] for obs in res['kis'][0]] assert max([d['DATE_BEG']['$date'] for d in desc]) < date.end.unix * 1000
date_end = [obs['description']['DATE_END']['$date'] for obs in res['kis'][0]] assert min([d['DATE_END']['$date'] for d in desc]) > date.start.unix * 1000
assert max(date_beg) < date.end.unix * 1000
assert min(date_end) > date.start.unix * 1000
date_beg = [obs['description']['DATE_BEG']['$date'] for obs in res['kis'][1]]
date_end = [obs['description']['DATE_END']['$date'] for obs in res['kis'][1]]
assert max(date_beg) < date.end.unix * 1000
assert min(date_end) > date.start.unix * 1000
files = Fido.fetch(res['kis'], binary=False) files = Fido.fetch(res['kis'][:, :100], binary=False)
assert len(files.errors) == 0 assert len(files.errors) == 0
assert len(files) >= 1 assert len(files) == 101
assert files[0].endswith('.json') assert files[0].endswith('.json')
for filepath in files: for filepath in files:
assert dirnames(filepath)[-1] in dirs
meta = json.load(open(filepath)) meta = json.load(open(filepath))
assert meta['_id']['$oid'] == os.path.splitext(os.path.basename(filepath))[0] assert meta['_id']['$oid'] == os.path.splitext(os.path.basename(filepath))[0]
assert date.start.iso[:10] in meta['metadata']['header']['DATE-OBS'] assert date.start.iso[:10] in meta['metadata']['header']['DATE-OBS']
assert Time(meta['metadata']['header']['DATE-OBS']).mjd < date.end.mjd + 600 assert Time(meta['metadata']['header']['DATE-OBS']).mjd < date.end.mjd + 600
files = Fido.fetch(res['kis'], binary=True) files = Fido.fetch(res['kis'][:, :10], binary=True)
assert files[0].endswith('.fits') assert files[0].endswith('.fits')
for filepath in files: for filepath in files:
hdulist = fits.open(filepath) hdulist = fits.open(filepath)
assert hdulist[0].header['TELESCOP'] in ('GREGOR', 'VTT') assert hdulist[0].header['TELESCOP'] in ('GREGOR', 'VTT')
assert hdulist[0].header['INSTRUME'] in ('GRIS', 'LARS') assert hdulist[0].header['INSTRUME'] in ('GRIS', 'LARS')
assert hdulist[0].header['EXTNAME'] in dirnames(filepath)[-1]
assert date.start.iso[:10] in hdulist[0].header['DATE-OBS'] assert date.start.iso[:10] in hdulist[0].header['DATE-OBS']
assert Time(hdulist[0].header['DATE-OBS']).mjd < date.end.mjd + 600 assert Time(hdulist[0].header['DATE-OBS']).mjd < date.end.mjd + 600
assert Time(hdulist[0].header['DATE-OBS']).mjd > date.start.mjd - 600 assert Time(hdulist[0].header['DATE-OBS']).mjd > date.start.mjd - 600
hdulist.close() hdulist.close()
assert len(files) == 301 assert len(files) == 11
@pytest.mark.parametrize("query", ((a.Instrument("GRIS") & a.Level(3)), @pytest.mark.parametrize("query", ((a.Instrument("GRIS") & a.Level(3)),
...@@ -449,7 +474,7 @@ def test_full_range(client): ...@@ -449,7 +474,7 @@ def test_full_range(client):
assert client._can_handle_query(query) assert client._can_handle_query(query)
if HAS_DOCKERTEST: if HAS_DOCKERTEST:
res = client.search(query) res = client.search(query)
assert len(res) == 1 assert len(res) == 105
assert 'description' in res.colnames assert 'description' in res.colnames
else: else:
with pytest.raises(URLError, match=rf"{_dockerexc('GRIS')}" with pytest.raises(URLError, match=rf"{_dockerexc('GRIS')}"
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment