Skip to content

GitLab

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
Show whitespace changes
Inline Side-by-side
sdc/client.py
View file @ 91d5e0eb
# -*- coding: utf-8 -*-
import warnings
import os
import copy
import warnings
import astropy.units as u
from astropy.time import Time
......@@ -274,7 +275,13 @@ class KISClient(BaseClient):
try:
response = urllib.request.urlopen(full_url)
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:
raise URLError(f'Unable to execute search "{full_url}": {exc}. Confirm that '
f'RESTHeart is running on {self._BASE_URL} and connected.')
......@@ -310,12 +317,13 @@ class KISClient(BaseClient):
for row in query_results:
inst = row['description']['INSTRUMENT']
rowpath = os.path.join(path, row['_id']['$oid'])
for l1_data in row['links']['l1_data']:
oid = l1_data['$oid']
filename = f"{oid}.{ext}"
oid = row['links']['l1_data']['$oid']
# Content-Disposition header default is "{row['_id']['$oid']}/{oid}.{ext}" (no '.json').
# rowpath = row['_id']['$oid']
filepath = os.path.join(row['description']['OBS_NAME'], f"{oid}.{ext}")
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
def _can_handle_query(cls, *query, hasinstr=False):
......
sdc/tests/test_client.py
View file @ 91d5e0eb
......@@ -2,6 +2,7 @@ import pytest
import urllib.request
from urllib.error import HTTPError, URLError
from pathlib import Path
import os
import json
import parfive
......@@ -9,6 +10,7 @@ import parfive
import astropy.units as u
from astropy.io import fits
from astropy.time import Time
import sunpy
from sunpy.net import Fido
from sunpy.net import attrs as a
from sunpy.net.base_client import QueryResponseTable
......@@ -31,6 +33,10 @@ except(HTTPError, URLError):
HAS_DOCKERTEST = False
def dirnames(path):
return os.path.dirname(path).split(os.path.sep)
@pytest.fixture
def client():
return KISClient()
......@@ -62,7 +68,7 @@ def test_docker(client):
res = client.search(a.Instrument("GRIS") & a.sdc.ObsName('gris_20140426_000'))
assert isinstance(res, QueryResponseTable)
assert len(res) == 1
assert len(res) == 105
description = res[0].get('description')
assert len(description) == 34
assert description['INSTRUMENT'] == 'gris'
......@@ -71,7 +77,7 @@ def test_docker(client):
assert description['DATE_BEG']['$date'] == 1398505619000
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
for oid in file_ids[0], file_ids[104]:
meta = json.loads(urllib.request.urlopen(f"{_BASE_URL}gris_l1_data.files/{oid}").read())
......@@ -87,8 +93,8 @@ def test_docker(client):
rowpath = f"{res[0]['_id']['$oid']}"
binfile = ''
ext = 'json'
for i, l1_data in enumerate(res[0]['links']['l1_data'][:10]):
oid = l1_data['$oid']
for i, ld in enumerate(res[:10]):
oid = ld['links']['l1_data']['$oid']
filename = f"{oid}.{ext}"
url = f"{_BASE_URL}{inst}_l1_data.files/{oid}{binfile}"
assert url == f"{_BASE_URL}gris_l1_data.files/{file_ids[i]}"
......@@ -96,15 +102,15 @@ def test_docker(client):
binfile = '/binary'
ext = 'fits'
for l1_data in res[0]['links']['l1_data'][:2]:
oid = l1_data['$oid']
for ld in res[:2]:
oid = ld['links']['l1_data']['$oid']
filename = f"{oid}.{ext}"
url = f"{_BASE_URL}{inst}_l1_data.files/{oid}{binfile}"
downloader.enqueue_file(url, filename=os.path.join(rowpath, filename), max_splits=1)
assert downloader.queued_downloads == 12
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)
files = downloader.download()
......@@ -159,7 +165,7 @@ def test_search(client):
if HAS_DOCKERTEST:
res = client.search(query)
assert isinstance(res, QueryResponseTable)
assert len(res) > 50
assert len(res) == 1
assert 'description' in res.colnames
else:
with pytest.raises(URLError, match=rf"{_dockerexc('bbi')}"
......@@ -173,7 +179,7 @@ def test_search(client):
assert client._can_handle_query(query)
if HAS_DOCKERTEST:
res = client.search(query)
assert len(res) == 100
assert len(res) == 1
assert 'description' in res.colnames
else:
with pytest.raises(URLError, match=rf"{_dockerexc('lars')}"
......@@ -185,7 +191,7 @@ def test_search(client):
assert client._can_handle_query(*query)
if HAS_DOCKERTEST:
res = client.search(*query)
assert len(res) == 100
assert len(res) == 1
assert 'description' in res.colnames
else:
with pytest.raises(URLError, match=rf"{_dockerexc('lars')}"
......@@ -197,7 +203,7 @@ def test_search(client):
assert client._can_handle_query(query)
if HAS_DOCKERTEST:
res = client.search(query)
assert len(res) == 200
assert len(res) == 362
assert 'description' in res.colnames
assert 'THETA' in res[0]['description']
theta = [obs['description']['THETA'] for obs in res]
......@@ -213,7 +219,7 @@ def test_search(client):
query = a.Instrument("LARS") | a.Instrument("GRIS"), a.sdc.Theta(85*u.deg, 3000*u.arcmin)
if HAS_DOCKERTEST:
res = client.search(*query)
assert len(res) == 200
assert len(res) == 163
assert 'description' in res.colnames
assert 'THETA' in res[0]['description']
theta = [obs['description']['THETA'] for obs in res]
......@@ -231,12 +237,12 @@ def test_fido_search():
two_inst = (a.Instrument("LARS") | a.Instrument("GRIS"))
if HAS_DOCKERTEST:
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']]
assert (min(theta) >= 50) & (max(theta) <= 80)
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']]
assert (min(theta) >= 50) & (max(theta) <= 80)
......@@ -250,7 +256,8 @@ def test_fido_search():
res = Fido.search(two_inst, a.sdc.Theta(50*u.deg, 80*u.deg))
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'
theta = [obs['description']['THETA'] for obs in res['kis'][0]]
assert (min(theta) >= 50) & (max(theta) <= 80)
......@@ -260,7 +267,7 @@ def test_fido_search():
date = a.Time("2016/08/26 16:25", "2016/08/26 16:45")
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'
else:
with pytest.raises(URLError, match=rf"{_dockerexc('gris')}"
......@@ -284,20 +291,26 @@ def test_fido_fetch():
date = a.Time("2017/05/22 08:45", "2017/05/22 08:55")
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'])
assert len(files) == 45
for filepath in files:
assert len(files) == 10
for i, filepath in enumerate(files):
assert dirnames(filepath)[-1] == desc[::-1][i]['OBS_NAME']
meta = json.load(open(filepath))
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 < 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)
assert os.path.dirname(files[0]).split(os.path.sep)[-1] == res['kis'][0]['_id']['$oid']
assert len(files) == 10
for filepath in files:
files = Fido.fetch(res['kis'][:3], binary=True)
assert len(files) == 3
for i, filepath in enumerate(files):
assert dirnames(filepath)[-1] == desc[::-1][i]['OBS_NAME']
hdulist = fits.open(filepath)
assert hdulist[0].header.get('TELESCOP') == 'GREGOR'
assert hdulist[0].header.get('INSTRUME') == 'BBI'
......@@ -320,13 +333,13 @@ def test_fido_fetch_2():
else:
res = Fido.search((a.Instrument("GRIS") | a.Instrument("LARS")) & date)
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_END']['$date'] > date.start.unix * 1000
files = Fido.fetch(res['kis'], binary=False)
assert len(files) == 400
files = Fido.fetch(res['kis'][0, :100], binary=False)
assert len(files) == 100
assert files[0].endswith('.json')
for filepath in files:
meta = json.load(open(filepath))
......@@ -334,9 +347,21 @@ def test_fido_fetch_2():
assert date.start.isot[:12] in 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")
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_END':{{'.gte':{{'.date':'{date.start.isot}'}}"):
res = Fido.search((a.Instrument("GRIS") | a.Instrument("LARS")) & date)
......@@ -344,39 +369,39 @@ def test_fido_fetch_2():
res = Fido.search((a.Instrument("GRIS") | a.Instrument("LARS")) & date)
assert len(res['kis']) == 2
assert len(res['kis'][0]) == 300
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]]
date_end = [obs['description']['DATE_END']['$date'] for obs in res['kis'][0]]
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
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'], binary=False)
files = Fido.fetch(res['kis'][:, :100], binary=False)
assert len(files.errors) == 0
assert len(files) >= 1
assert len(files) == 101
assert files[0].endswith('.json')
for filepath in files:
assert dirnames(filepath)[-1] in dirs
meta = json.load(open(filepath))
assert meta['_id']['$oid'] == os.path.splitext(os.path.basename(filepath))[0]
assert date.start.iso[:10] in meta['metadata']['header']['DATE-OBS']
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')
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 hdulist[0].header['EXTNAME'] in dirnames(filepath)[-1]
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()
assert len(files) == 301
assert len(files) == 11
@pytest.mark.parametrize("query", ((a.Instrument("GRIS") & a.Level(3)),
......@@ -449,7 +474,7 @@ def test_full_range(client):
assert client._can_handle_query(query)
if HAS_DOCKERTEST:
res = client.search(query)
assert len(res) == 1
assert len(res) == 105
assert 'description' in res.colnames
else:
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