1
|
|
|
# -*- coding: utf-8 -*- |
2
|
|
|
# vim:fileencoding=utf-8 |
3
|
|
|
# |
4
|
|
|
# Copyright (c) 2022 Stefan Bender |
5
|
|
|
# |
6
|
|
|
# This module is part of pyspaceweather. |
7
|
|
|
# pyspaceweather is free software: you can redistribute it or modify |
8
|
|
|
# it under the terms of the GNU General Public License as published |
9
|
|
|
# by the Free Software Foundation, version 2. |
10
|
|
|
# See accompanying COPYING.GPLv2 file or http://www.gnu.org/licenses/gpl-2.0.html. |
11
|
|
|
"""OMNI data read tests |
12
|
|
|
""" |
13
|
|
|
import os |
14
|
|
|
import requests |
15
|
|
|
|
16
|
|
|
import numpy as np |
17
|
|
|
import pandas as pd |
18
|
|
|
|
19
|
|
|
import pytest |
20
|
|
|
|
21
|
|
|
from spaceweather import cache_omnie, omnie_hourly, sw_daily |
22
|
|
|
from spaceweather.omni import OMNI_URL_BASE, OMNI_PREFIX, OMNI_EXT |
23
|
|
|
|
24
|
|
|
_TEST_YEAR = 2012 |
25
|
|
|
_TEST_FILE = "{0}_{1:04d}.{2}".format(OMNI_PREFIX, _TEST_YEAR, OMNI_EXT) |
26
|
|
|
_TEST_URL = os.path.join(OMNI_URL_BASE, _TEST_FILE) |
27
|
|
|
|
28
|
|
|
|
29
|
|
|
def test_cache(mocker, tmpdir): |
30
|
|
|
mocker.patch("requests.get") |
31
|
|
|
tmpdir = str(tmpdir) |
32
|
|
|
cache_omnie(year=_TEST_YEAR, local_path=tmpdir) |
33
|
|
|
requests.get.assert_called_once_with(_TEST_URL, stream=True) |
34
|
|
|
|
35
|
|
|
|
36
|
|
|
def test_auto_update(mocker, tmpdir): |
37
|
|
|
# test with non-existent file |
38
|
|
|
mocker.patch("requests.get") |
39
|
|
|
tmpdir = str(tmpdir) |
40
|
|
|
with pytest.raises(IOError): |
41
|
|
|
with pytest.warns(UserWarning): |
42
|
|
|
omnie_hourly(year=_TEST_YEAR, cache=True, local_path=tmpdir) |
43
|
|
|
requests.get.assert_called_once_with(_TEST_URL, stream=True) |
44
|
|
|
|
45
|
|
|
|
46
|
|
|
def test_not_avail(mocker, tmpdir): |
47
|
|
|
# test with non-existent file |
48
|
|
|
tmpdir = str(tmpdir) |
49
|
|
|
with pytest.raises(IOError): |
50
|
|
|
with pytest.warns(UserWarning): |
51
|
|
|
omnie_hourly(year=_TEST_YEAR, cache=False, local_path=tmpdir) |
52
|
|
|
|
53
|
|
|
|
54
|
|
|
@pytest.mark.parametrize("hour", range(0, 24, 3)) |
55
|
|
|
@pytest.mark.parametrize("index", ["Ap", "Kp"]) |
56
|
|
|
def test_hourly(hour, index): |
57
|
|
|
df1 = omnie_hourly(2000) |
58
|
|
|
df2 = sw_daily() |
59
|
|
|
ind_name = "{0}{1}".format(index, hour) |
60
|
|
|
df1_ind = df1[df1["hour"] == hour][index] |
61
|
|
|
df2_ind = df2[ind_name].loc[df1_ind.index.date].rename(index) |
62
|
|
|
df2_ind.index = df1_ind.index |
63
|
|
|
pd.testing.assert_series_equal(df1_ind, df2_ind) |
64
|
|
|
|
65
|
|
|
|
66
|
|
View Code Duplication |
@pytest.mark.parametrize( |
|
|
|
|
67
|
|
|
"name, result", |
68
|
|
|
[ |
69
|
|
|
("Ap", np.array([56, 39, 27, 18, 32, 15, 32, 22])), |
70
|
|
|
("Kp", np.array([5.3, 4.7, 4.0, 3.3, 4.3, 3.0, 4.3, 3.7])), |
71
|
|
|
] |
72
|
|
|
) |
73
|
|
|
def test_3hourly_index(name, result): |
74
|
|
|
df = omnie_hourly(2000) |
75
|
|
|
np.testing.assert_allclose( |
76
|
|
|
df.loc[ |
77
|
|
|
pd.date_range( |
78
|
|
|
"2000-01-01 00:00", "2000-01-01 23:00", freq="3h" |
79
|
|
|
) |
80
|
|
|
][name].values, |
81
|
|
|
result, |
82
|
|
|
rtol=1e-12, |
83
|
|
|
) |
84
|
|
|
|