Controller_test.py

#! /usr/bin/env python
from geomagio import Controller, TimeseriesFactory
from geomagio.algorithm import Algorithm

# needed to read outputs generated by Controller and test data
from geomagio.iaga2002 import IAGA2002Factory

# needed to emulate geomag.py script
from geomagio.Controller import _main, get_previous_interval, parse_args

# needed to copy SqDistAlgorithm statefile
from shutil import copy

# needed to determine a valid (and writable) temp folder
from tempfile import gettempdir

from numpy.testing import assert_allclose, assert_equal
from obspy.core import UTCDateTime


def test_controller():
    """Controller_test.test_controller()

    instantiate the controller, make certain the factories and algorithms
    are set
    """
    inputfactory = TimeseriesFactory()
    outputfactory = TimeseriesFactory()
    algorithm = Algorithm()
    controller = Controller(inputfactory, outputfactory, algorithm)
    assert_equal(isinstance(controller._inputFactory, TimeseriesFactory), True)
    assert_equal(isinstance(controller._outputFactory, TimeseriesFactory), True)
    assert_equal(isinstance(controller._algorithm, Algorithm), True)


def test_controller_input_url():
    """Controller_test.test_controller_input_url()

    Test `--input-url` option in controller with both a URL template, and
    and a simple, single file url. The latter was especially problematic after
    switching to Py3 because BytesIO choked on data read in as a string.
    NOTE: the bug this was designed to catch was TypeError, so a "pass" is
          simply to read in the data file. There is no need to compare it
          with anything.
    """
    # define folder for testing
    tmp_dir = gettempdir()

    # TEST 1 - read in interval using the --input-url option set to a
    # URL template (i.e., it has a '{' in the string)

    # create list of string command line arguments
    fake_argv = [
        "--input",
        "iaga2002",
        "--input-url",
        "file://etc/controller/{obs}{date:%Y%m%d}_XYZF_{t}{i}.{i}",
        "--observatory",
        "BOU",
        "--inchannels",
        "X",
        "Y",
        "Z",
        "F",
        "--interval",
        "minute",
        "--type",
        "variation",
        "--output",
        "iaga2002",
        "--output-url",
        "file://" + tmp_dir + "/{obs}{date:%Y%m%d}_XYZF_{t}{i}.{i}",
    ]

    # parse arguments and create initial args object
    args = parse_args(fake_argv)

    starttime1 = args.starttime = UTCDateTime("2018-10-24T00:00:00Z")
    endtime1 = args.endtime = UTCDateTime("2018-10-24T00:19:00Z")
    _main(args)

    # TEST 2 - read in interval using the --input-url option with no
    # URL template (i.e., no '{' in the string, just a single filename)
    # create list of string command line arguments
    fake_argv = [
        "--input",
        "iaga2002",
        "--input-url",
        "file://etc/controller/bou20181024_XYZF_vmin.min",
        "--observatory",
        "BOU",
        "--inchannels",
        "X",
        "Y",
        "Z",
        "F",
        "--interval",
        "minute",
        "--type",
        "variation",
        "--output",
        "iaga2002",
        "--output-url",
        "file://" + tmp_dir + "/bou20181024_XYZF_noURL_vmin.min",
    ]

    # parse arguments and create initial args object
    args = parse_args(fake_argv)

    starttime1 = args.starttime = UTCDateTime("2018-10-24T00:00:00Z")
    endtime1 = args.endtime = UTCDateTime("2018-10-24T00:19:00Z")
    _main(args)


def test_controller_update_sqdist():
    """Controller_test.test_controller_update_sqdist().

    This is an end-to-end test of the Controller, more-or-less how it would be
    invoked via the geomag.py command line script. We specifically test the
    Controller's run() logic using the SqDistAlgoritm and carefully
    constructed inputs since this is one of the most complicated anticipated
    use-cases. Some liberties have been taken to avoid repeatedly parsing all
    arguments or reloading the interpreter.

    This test also takes advantage of the fact that args.realtime is processed
    at the end of main() before _main() is called.  This test explicitly
    sets starttime, endtime, and realtime argument values to override what
    may otherwise be expected during normal command line operations.
    """
    # define folder for testing
    tmp_dir = gettempdir()

    # create list of string command line arguments
    fake_argv = [
        "--input",
        "iaga2002",
        "--input-url",
        "file://etc/controller/{obs}{date:%Y%m%d}_XYZF_{t}{i}.{i}",
        "--observatory",
        "BOU",
        "--algorithm",
        "sqdist",
        "--sqdist-m",
        "1440",
        "--sqdist-alpha",
        "2.3148e-5",
        "--sqdist-gamma",
        "3.3333e-2",
        "--sqdist-smooth",
        "180",
        "--inchannels",
        "X",
        "Y",
        "Z",
        "F",
        "--interval",
        "minute",
        "--rename-output-channel",
        "H_Dist",
        "MDT",
        "--rename-output-channel",
        "H_SQ",
        "MSQ",
        "--rename-output-channel",
        "H_SV",
        "MSV",
        "--rename-output-channel",
        "H_Sigma",
        "MSS",
        "--outchannels",
        "MDT",
        "MSQ",
        "MSV",
        "MSS",
        "--sqdist-mag",
        "--sqdist-statefile",
        tmp_dir + "/sqdistBOU_h_state.json",
        "--type",
        "variation",
        "--output",
        "iaga2002",
        "--output-url",
        "file://" + tmp_dir + "/{obs}{date:%Y%m%d}_DQVS_{t}{i}.{i}",
        "--realtime",
        "600",
    ]
    # parse arguments and create initial args object
    args = parse_args(fake_argv)

    # read in test and latest output and compare
    actual_factory = IAGA2002Factory(
        urlTemplate=("file://" + tmp_dir + "/{obs}{date:%Y%m%d}_DQVS_{t}{i}.{i}"),
        urlInterval=86400,
        observatory="BOU",
        channels=["MDT", "MSQ", "MSV", "MSS"],
    )
    expected_factory = IAGA2002Factory(
        urlTemplate="url template, individual tests change the template below",
        urlInterval=86400,
        observatory="BOU",
        channels=["MDT", "MSQ", "MSV", "MSS"],
    )

    # setup test data
    # copy SqDistAlgorithm statefile and empty DQVS output file to tmp folder
    copy("etc/controller/sqdistBOU_h_state.json", tmp_dir)
    copy(
        "etc/controller/bou20181024_DQVS_test0_vmin.min",
        tmp_dir + "/bou20181024_DQVS_vmin.min",
    )

    # TEST 1 - include a gap at end that is less than realtime (10 minutes),
    # expect sqdist not to project SQ/SV/SS
    starttime1 = args.starttime = UTCDateTime("2018-10-24T00:00:00Z")
    endtime1 = args.endtime = UTCDateTime("2018-10-24T00:19:00Z")
    _main(args)
    # compare results
    actual = actual_factory.get_timeseries(starttime=starttime1, endtime=endtime1)
    expected_factory.urlTemplate = (
        "file://etc/controller/{obs}{date:%Y%m%d}_DQVS_test1_{t}{i}.{i}"
    )
    expected = expected_factory.get_timeseries(starttime=starttime1, endtime=endtime1)
    assert_allclose(actual, expected)

    # TEST 2 - start after next_starttime (00:10),
    # expect SQDist to project sq/sv/ss values over gap,
    # then process until last gap starting at 00:38
    args.startime = UTCDateTime("2018-10-24T00:20:00Z")
    endtime2 = args.endtime = UTCDateTime("2018-10-24T00:39:00Z")
    _main(args)
    # compare results
    actual = actual_factory.get_timeseries(starttime=starttime1, endtime=endtime2)
    expected_factory.urlTemplate = (
        "file://etc/controller/{obs}{date:%Y%m%d}_DQVS_test2_{t}{i}.{i}"
    )
    expected = expected_factory.get_timeseries(starttime=starttime1, endtime=endtime2)
    assert_allclose(actual, expected)

    # TEST 3 - start after next_starttime (00:38),
    # expect SQDist to project over gap,
    # then process until last gap starting at 00:58
    args.starttime = UTCDateTime("2018-10-24T00:40:00Z")
    endtime3 = args.endtime = UTCDateTime("2018-10-24T00:59:00Z")
    _main(args)
    # compare results
    actual = actual_factory.get_timeseries(starttime=starttime1, endtime=endtime3)
    expected_factory.urlTemplate = (
        "file://etc/controller/{obs}{date:%Y%m%d}_DQVS_test3_{t}{i}.{i}"
    )
    expected = expected_factory.get_timeseries(starttime=starttime1, endtime=endtime3)
    assert_allclose(actual, expected)

    # TEST 4 - start after next_starttime (00:58),
    # exptect SQDist to project over gap,
    # then process until last gap starting at 01:16
    args.starttime = UTCDateTime("2018-10-24T01:00:00Z")
    endtime4 = args.endtime = UTCDateTime("2018-10-24T01:19:00Z")
    _main(args)
    # compare results
    actual = actual_factory.get_timeseries(starttime=starttime1, endtime=endtime4)
    expected_factory.urlTemplate = (
        "file://etc/controller/{obs}{date:%Y%m%d}_DQVS_test4_{t}{i}.{i}"
    )
    expected = expected_factory.get_timeseries(starttime=starttime1, endtime=endtime4)
    assert_allclose(actual, expected)

    # TEST 5 - start after next_starttime (01:16),
    # expect SQDist to project until beginning of realtime gap,
    # starting at 01:30 (01:39 - 600 seconds)
    args.starttime = UTCDateTime("2018-10-24T01:20:00Z")
    endtime5 = args.endtime = UTCDateTime("2018-10-24T01:39:00Z")
    _main(args)
    # compare results
    actual = actual_factory.get_timeseries(starttime=starttime1, endtime=endtime5)
    expected_factory.urlTemplate = (
        "file://etc/controller/{obs}{date:%Y%m%d}_DQVS_test5_{t}{i}.{i}"
    )
    expected = expected_factory.get_timeseries(starttime=starttime1, endtime=endtime5)
    assert_allclose(actual, expected)

    # TEST 6 - set starttime before next_starttime (which is 01:30)
    # expect sqdist to pick up where it left off
    args.starttime = UTCDateTime("2018-10-24T01:20:00Z")
    endtime6 = args.endtime = UTCDateTime("2018-10-24T01:59:00Z")
    _main(args)
    # compare results
    actual = actual_factory.get_timeseries(starttime=starttime1, endtime=endtime6)
    expected_factory.urlTemplate = (
        "file://etc/controller/{obs}{date:%Y%m%d}_DQVS_test6_{t}{i}.{i}"
    )
    expected = expected_factory.get_timeseries(starttime=starttime1, endtime=endtime6)
    assert_allclose(actual, expected)


def test_get_previous_interval():
    """Test get_previous_interval produces stable interval sizes"""
    # initial interval is one hour
    start = UTCDateTime("2022-01-05T00:00:00Z")
    end = UTCDateTime("2022-01-05T01:00:00Z")
    # previous interval starts at beginning of previous hour
    previous = get_previous_interval(start=start, end=end)
    assert previous == (
        UTCDateTime("2022-01-04T23:00:00"),
        UTCDateTime("2022-01-04T23:59:59.999999Z"),
    )
    # previous interval still starts at beginning of previous hour
    # even though interval is one microsecond smaller
    previous = get_previous_interval(*previous)
    assert previous == (
        UTCDateTime("2022-01-04T22:00:00"),
        UTCDateTime("2022-01-04T22:59:59.999999Z"),
    )