DataApiQuery_test.py

import datetime
import os
from numpy.testing import assert_equal
from obspy import UTCDateTime

from geomagio.api.ws.DataApiQuery import (
    DataApiQuery,
    SamplingPeriod,
    DataType,
    OutputFormat,
    DataHost,
)


def test_DataApiQuery_defaults():
    query = DataApiQuery(id="BOU")

    now = datetime.datetime.now(tz=datetime.timezone.utc)
    expected_start_time = UTCDateTime(year=now.year, month=now.month, day=now.day)
    expected_endtime = expected_start_time + (86400 - 0.001)

    assert_equal(query.id, "BOU")
    assert_equal(query.starttime, expected_start_time)
    assert_equal(query.endtime, expected_endtime)
    assert_equal(query.elements, ["X", "Y", "Z", "F"])
    assert_equal(query.sampling_period, SamplingPeriod.MINUTE)
    assert_equal(query.data_type, DataType.VARIATION)
    assert_equal(query.format, OutputFormat.IAGA2002)
    # assumes the env var DATA_HOST is not set
    assert_equal(query.data_host.value, "edgecwb.usgs.gov")


def test_DataApiQuery_starttime_is_none():
    query = DataApiQuery(id="BOU", starttime=None)

    now = datetime.datetime.now(tz=datetime.timezone.utc)
    expected_start_time = UTCDateTime(year=now.year, month=now.month, day=now.day)
    expected_endtime = expected_start_time + (86400 - 0.001)

    assert_equal(query.id, "BOU")
    assert_equal(query.starttime, expected_start_time)
    assert_equal(query.endtime, expected_endtime)
    assert_equal(query.elements, ["X", "Y", "Z", "F"])
    assert_equal(query.sampling_period, SamplingPeriod.MINUTE)
    assert_equal(query.data_type, DataType.VARIATION)
    assert_equal(query.format, OutputFormat.IAGA2002)
    # assumes the env var DATA_HOST is not set
    assert_equal(query.data_host.value, "edgecwb.usgs.gov")


def test_DataApiQuery_valid():
    query = DataApiQuery(
        id="DED",
        starttime="2024-09-01T00:00:01",
        endtime="2024-09-01T01:00:01",
        elements=["F"],
        sampling_period=1,
        data_type="adjusted",
        format="json",
        data_host="cwbpub.cr.usgs.gov",
    )

    assert_equal(query.id, "DED")
    assert_equal(query.starttime, UTCDateTime("2024-09-01T00:00:01"))
    assert_equal(query.endtime, UTCDateTime("2024-09-01T01:00:01"))
    assert_equal(query.elements, ["F"])
    assert_equal(query.sampling_period, SamplingPeriod.SECOND)
    assert_equal(query.data_type, "adjusted")
    assert_equal(query.format, "json")
    assert_equal(query.data_host, "cwbpub.cr.usgs.gov")


def test_DataApiQuery_no_id():
    query = None
    try:
        query = DataApiQuery()
    except Exception as e:
        err = e.errors()
        assert "Field required" == err[0]["msg"]

    assert_equal(query, None)


def test_DataApiQuery_invalid_id():
    query = None
    try:
        query = DataApiQuery(id="YMCA")
    except Exception as e:
        err = e.errors()
        assert "Value error, Bad observatory id 'YMCA'." in err[0]["msg"]

    assert_equal(query, None)


def test_DataApiQuery_default_endtime():
    query = DataApiQuery(id="BOU", starttime="2024-11-01T00:00:01")

    assert_equal(query.id, "BOU")
    assert_equal(query.starttime, UTCDateTime("2024-11-01T00:00:01"))
    # endtime is 1 day after start time
    assert_equal(query.endtime, UTCDateTime("2024-11-02T00:00:00.999"))
    assert_equal(query.elements, ["X", "Y", "Z", "F"])
    assert_equal(query.sampling_period, SamplingPeriod.MINUTE)
    assert_equal(query.data_type, DataType.VARIATION)
    assert_equal(query.format, OutputFormat.IAGA2002)
    assert_equal(query.data_host, DataHost.DEFAULT)


def test_DataApiQuery_default_only_endtime():
    # test a valid endtime that is after default starttime
    hour_later = datetime.datetime.now() + datetime.timedelta(hours=1)

    query = DataApiQuery(id="BOU", endtime=hour_later)

    assert_equal(query.id, "BOU")

    now = datetime.datetime.now(tz=datetime.timezone.utc)
    expected_start_time = UTCDateTime(year=now.year, month=now.month, day=now.day)
    # starttime is at the beginning of today
    assert_equal(query.starttime, expected_start_time)
    # endtime is still what the user set
    assert_equal(query.endtime, hour_later)

    assert_equal(query.elements, ["X", "Y", "Z", "F"])
    assert_equal(query.sampling_period, SamplingPeriod.MINUTE)
    assert_equal(query.data_type, DataType.VARIATION)
    assert_equal(query.format, OutputFormat.IAGA2002)
    assert_equal(query.data_host, DataHost.DEFAULT)


def test_DataApiQuery_default_only_invalid_endtime():
    # test an invalid endtime that is before default starttime
    query = None
    try:
        query = DataApiQuery(id="BOU", endtime="2024-09-01T00:00:01")
    except Exception as e:
        err = e.errors()
        assert "Value error, Starttime must be before endtime." == err[0]["msg"]

    assert_equal(query, None)


def test_DataApiQuery_starttime_after_endtime():
    query = None
    try:
        query = DataApiQuery(
            id="BOU", starttime="2024-11-01T00:00:01", endtime="2024-10-01T00:00:01"
        )
    except Exception as e:
        err = e.errors()
        assert "Value error, Starttime must be before endtime." == err[0]["msg"]

    assert_equal(query, None)


def test_DataApiQuery_invalid_starttime():
    query = None
    try:
        query = DataApiQuery(
            id="BOU",
            starttime="November 8th 2024 12:00pm",
        )
    except Exception as e:
        err = e.errors()
        assert "Input should be an instance of UTCDateTime" == err[0]["msg"]

    assert_equal(query, None)


def test_DataApiQuery_invalid_endtime():
    query = None
    try:
        query = DataApiQuery(
            id="BOU",
            starttime="2024-11-01T00:00:01",
            endtime="Add one day",
        )
    except Exception as e:
        err = e.errors()
        assert "Input should be an instance of UTCDateTime" == err[0]["msg"]

    assert_equal(query, None)


def test_DataApiQuery_invalid_list_elements():
    query = None
    try:
        query = DataApiQuery(id="BOU", elements="XYZF")
    except Exception as e:
        err = e.errors()
        assert "Input should be a valid list" == err[0]["msg"]

    assert_equal(query, None)


def test_DataApiQuery_invalid_element():
    query = None
    try:
        query = DataApiQuery(id="BOU", elements=["A"])
    except Exception as e:
        err = e.errors()
        assert (
            "Value error, Bad element 'A'. Valid values are: U, V, W, H, E, X, Y, D, Z, F, G, DIST, E-E, E-N, SQ, SV, UK1, UK2, UK3, UK4."
            == err[0]["msg"]
        )

    assert_equal(query, None)


def test_DataApiQuery_invalid_elements():
    query = None
    try:
        query = DataApiQuery(id="BOU", elements=["Z", "A", "F"])
    except Exception as e:
        err = e.errors()
        assert (
            "Value error, Bad element 'A'. Valid values are: U, V, W, H, E, X, Y, D, Z, F, G, DIST, E-E, E-N, SQ, SV, UK1, UK2, UK3, UK4."
            == err[0]["msg"]
        )

    assert_equal(query, None)


def test_DataApiQuery_invalid_sampling_period():
    query = None
    try:
        query = DataApiQuery(id="BOU", sampling_period=100)
    except Exception as e:
        err = e.errors()
        assert "Input should be 0.1, 1.0, 60.0, 3600.0 or 86400.0" == err[0]["msg"]

    assert_equal(query, None)


def test_DataApiQuery_extra_fields():
    query = None
    try:
        query = DataApiQuery(id="BOU", foo="bar")
    except Exception as e:
        err = e.errors()
        assert "Extra inputs are not permitted" == err[0]["msg"]

    assert_equal(query, None)