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, 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_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)