test_sds.py

# -*- coding: utf-8 -*-
import imghdr
import inspect
import os
import re
import shutil
import tempfile
import unittest

import numpy as np

from obspy import UTCDateTime, Trace, Stream
from obspy.core.util.misc import TemporaryWorkingDirectory
from obspy.clients.filesystem.sds import SDS_FMTSTR, Client
from obspy.scripts.sds_html_report import main as sds_report


class TemporarySDSDirectory(object):
    """
    Handles creation and deletion of a temporary SDS directory structure.
    To be used with "with" statement.
    """
    sampling_rate = 0.1
    networks = ("AB", "CD")
    stations = ("XYZ", "ZZZ3")
    locations = ("", "00")
    channels = ("HHZ", "HHN", "HHE", "BHZ", "BHN", "BHE")

    def __init__(self, year, doy, time=None):
        """
        Set which day's midnight (00:00 hours) is used as a day break in the
        testing (to split the test data into two files).

        If `time` is specified it overrides `year` and `doy`.
        """
        if time:
            self.time = time
        else:
            self.time = UTCDateTime("%d-%03dT00:00:00" % (year, doy))
        delta = 1.0 / self.sampling_rate

        self.stream = Stream()
        for net in self.networks:
            for sta in self.stations:
                for loc in self.locations:
                    for cha in self.channels:
                        tr = Trace(
                            data=np.arange(100, dtype=np.int32),
                            header=dict(
                                network=net, station=sta, location=loc,
                                channel=cha, sampling_rate=self.sampling_rate,
                                starttime=self.time - 30 * delta))

                        # cut into two seamless traces
                        tr1 = tr.slice(endtime=self.time + 5 * delta)
                        tr2 = tr.slice(starttime=self.time + 6 * delta)
                        self.stream.append(tr1)
                        self.stream.append(tr2)

    def __enter__(self):
        self.old_dir = os.getcwd()
        self.tempdir = tempfile.mkdtemp(prefix='obspy-sdstest-')
        for tr_ in self.stream:
            t_ = tr_.stats.starttime
            full_path = SDS_FMTSTR.format(year=t_.year, doy=t_.julday,
                                          sds_type="D", **tr_.stats)
            full_path = os.path.join(self.tempdir, full_path)
            dirname, filename = os.path.split(full_path)
            if not os.path.isdir(dirname):
                os.makedirs(dirname)
            tr_.write(full_path, format="MSEED")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):  # @UnusedVariable
        os.chdir(self.old_dir)
        shutil.rmtree(self.tempdir)


class SDSTestCase(unittest.TestCase):
    """
    Test reading data from SDS file structure.
    """
    def setUp(self):
        self.data_dir = os.path.join(os.path.dirname(os.path.abspath(
            inspect.getfile(inspect.currentframe()))), "data")

    def test_read_from_sds(self):
        """
        Test reading data across year and day breaks from SDS directory
        structure. Also tests looking for data on the wrong side of a day
        break (data usually get written for some seconds into the previous or
        next file around midnight).
        """
        # test for one specific SEED ID, without wildcards
        seed_id = "AB.XYZ..HHZ"
        net, sta, loc, cha = seed_id.split(".")
        # use three different day breaks in the testing:
        #  - normal day break during one year
        #     (same directory, separate filenames)
        #  - day break at end of year
        #     (separate directories, separate filenames)
        #      - leap-year
        #      - non-leap-year
        for year, doy in ((2015, 123), (2015, 1), (2012, 1)):
            t = UTCDateTime("%d-%03dT00:00:00" % (year, doy))
            with TemporarySDSDirectory(year=year, doy=doy) as temp_sds:
                # normal test reading across the day break
                client = Client(temp_sds.tempdir)
                st = client.get_waveforms(net, sta, loc, cha, t - 20, t + 20)
                self.assertEqual(len(st), 1)
                self.assertEqual(st[0].stats.starttime, t - 20)
                self.assertEqual(st[0].stats.endtime, t + 20)
                self.assertEqual(len(st[0]), 5)
                # test merge option
                st = client.get_waveforms(net, sta, loc, cha, t - 200, t + 200,
                                          merge=False)
                self.assertEqual(len(st), 2)
                st = client.get_waveforms(net, sta, loc, cha, t - 200, t + 200,
                                          merge=None)
                self.assertEqual(len(st), 2)
                st = client.get_waveforms(net, sta, loc, cha, t - 200, t + 200,
                                          merge=0)
                self.assertEqual(len(st), 1)
                # test reading data from a single day file
                # (data is in the file where it's expected)
                st = client.get_waveforms(net, sta, loc, cha, t - 80, t - 30)
                self.assertEqual(len(st), 1)
                # test reading data from a single day file
                # (data is in the dayfile of the previous day)
                st = client.get_waveforms(net, sta, loc, cha, t + 20, t + 40)
                self.assertEqual(len(st), 1)
                # test that format autodetection with `format=None` works
                client = Client(temp_sds.tempdir, format=None)
                st = client.get_waveforms(net, sta, loc, cha, t - 200, t + 200)
                self.assertEqual(len(st), 1)

    def test_read_from_sds_with_wildcarded_seed_ids(self):
        """
        Test reading data with wildcarded SEED IDs.
        """
        year, doy = 2015, 1
        t = UTCDateTime("%d-%03dT00:00:00" % (year, doy))
        with TemporarySDSDirectory(year=year, doy=doy) as temp_sds:
            # test different wildcard combinations in SEED ID
            client = Client(temp_sds.tempdir)
            for wildcarded_seed_id, num_matching_ids in zip(
                    ("AB.ZZZ3..HH?", "AB.ZZZ3..HH*", "*.*..HHZ",
                     "*.*.*.HHZ", "*.*.*.*"),
                    (3, 3, 4, 8, 48)):
                net, sta, loc, cha = wildcarded_seed_id.split(".")
                st = client.get_waveforms(net, sta, loc, cha, t - 200, t + 200)
                self.assertEqual(len(st), num_matching_ids)
            # test with SDS type wildcards
            for type_wildcard in ("*", "?"):
                net, sta, loc, cha = wildcarded_seed_id.split(".")
                st = client.get_waveforms(net, sta, loc, cha, t - 200, t + 200)
                self.assertEqual(len(st), num_matching_ids)

    def test_sds_report(self):
        """
        Test command line script for generating SDS report html.

        Inherently that script uses many other routines like `_get_filenames`,
        `get_availability_percentage`, `_get_current_endtime`,
        `get_latency`, `has_data` and `get_all_stations`, so these should be
        sufficiently covered as well.
        """
        # generate some dummy SDS with data roughly 2-3 hours behind current
        # time
        t = UTCDateTime() - 2.5 * 3600
        with TemporarySDSDirectory(year=None, doy=None, time=t) as temp_sds, \
                TemporaryWorkingDirectory():
            # create the report
            output_basename = "sds_report"
            argv = [
                "-r={}".format(temp_sds.tempdir),
                "-o={}".format(os.path.join(os.curdir, output_basename)),
                "-l=", "-l=00", "-l=10", "-c=HHZ", "-c=BHZ", "-i=AB.XYZ..BHE",
                "--check-quality-days=1"]
            sds_report(argv)
            # do the testing
            output_basename_abspath = os.path.abspath(
                os.path.join(os.curdir, output_basename))
            file_html = output_basename_abspath + ".html"
            file_txt = output_basename_abspath + ".txt"
            file_png = output_basename_abspath + ".png"
            # check that output files exist
            for file_ in [file_html, file_txt, file_png]:
                self.assertTrue(os.path.isfile(file_))
            # check content of image file (just check it is a png file)
            self.assertEqual(imghdr.what(file_png), "png")
            # check content of stream info / data quality file
            expected_lines = [
                b"AB,XYZ,,BHE,831[0-9].[0-9]*?,0.007292,2",
                b"AB,XYZ,,HHZ,831[0-9].[0-9]*?,0.007292,2",
                b"AB,XYZ,00,HHZ,831[0-9].[0-9]*?,0.007292,2",
                b"AB,ZZZ3,,HHZ,831[0-9].[0-9]*?,0.007292,2",
                b"AB,ZZZ3,00,HHZ,831[0-9].[0-9]*?,0.007292,2",
                b"CD,XYZ,,HHZ,831[0-9].[0-9]*?,0.007292,2",
                b"CD,XYZ,00,HHZ,831[0-9].[0-9]*?,0.007292,2",
                b"CD,ZZZ3,,HHZ,831[0-9].[0-9]*?,0.007292,2",
                b"CD,ZZZ3,00,HHZ,831[0-9].[0-9]*?,0.007292,2"]
            with open(file_txt, "rb") as fh:
                got_lines = fh.readlines()
            for expected_line, got_line in zip(expected_lines, got_lines):
                self.assertIsNotNone(re.match(expected_line, got_line))
            # check content of html report
            with open(file_html, "rb") as fh:
                got_lines = fh.readlines()
            html_regex_file = os.path.join(self.data_dir, "sds_report.regex")
            with open(html_regex_file, "rb") as fh:
                regex_patterns = fh.readlines()
            failed = False  # XXX remove again
            for got, pattern in zip(got_lines, regex_patterns):
                match = re.match(pattern.strip(), got.strip())
                try:
                    self.assertIsNotNone(match)
                except AssertionError:
                    failed = True
                    print(pattern.strip())
                    print(got.strip())
            if failed:
                raise Exception

    def test_get_waveforms_bulk(self):
        """
        Test get_waveforms_bulk method.
        """
        year = 2015
        doy = 247
        t = UTCDateTime("%d-%03dT00:00:00" % (year, doy))
        with TemporarySDSDirectory(year=year, doy=doy) as temp_sds:
            chunks = [
                ["AB", "XYZ", "", "HHZ", t, t + 20],
                ["AB", "XYZ", "", "HHN", t + 20, t + 40],
                ["AB", "XYZ", "", "HHE", t + 40, t + 60],
                ["CD", "ZZZ3", "00", "BHZ", t + 60, t + 80],
                ["CD", "ZZZ3", "00", "BHN", t + 80, t + 100],
                ["CD", "ZZZ3", "00", "BHE", t + 120, t + 140]
            ]
            client = Client(temp_sds.tempdir)
            st = client.get_waveforms_bulk(chunks)
            for _i in range(6):
                if _i <= 2:
                    self.assertEqual(st[_i].stats.network, "AB")
                    self.assertEqual(st[_i].stats.station, "XYZ")
                    self.assertEqual(st[_i].stats.location, "")
                elif _i >= 2:
                    self.assertEqual(st[_i].stats.network, "CD")
                    self.assertEqual(st[_i].stats.station, "ZZZ3")
                    self.assertEqual(st[_i].stats.location, "00")
            self.assertEqual(st[0].stats.channel, "HHZ")
            self.assertEqual(st[1].stats.channel, "HHN")
            self.assertEqual(st[2].stats.channel, "HHE")
            self.assertEqual(st[3].stats.channel, "BHZ")
            self.assertEqual(st[4].stats.channel, "BHN")
            self.assertEqual(st[5].stats.channel, "BHE")

    def test_get_all_stations_and_nslc(self):
        """
        Test `get_all_stations` and `get_all_nslc` methods
        """
        # generate dummy SDS
        t = UTCDateTime()
        with TemporarySDSDirectory(year=None, doy=None, time=t) as temp_sds:
            client = Client(temp_sds.tempdir)
            expected_netsta = sorted([
                (net, sta)
                for net in temp_sds.networks
                for sta in temp_sds.stations])
            got_netsta = client.get_all_stations()
            self.assertEqual(expected_netsta, got_netsta)
            expected_nslc = sorted([
                (net, sta, loc, cha)
                for net in temp_sds.networks
                for sta in temp_sds.stations
                for loc in temp_sds.locations
                for cha in temp_sds.channels])
            got_nslc = client.get_all_nslc()
            self.assertEqual(expected_nslc, got_nslc)
            got_nslc = client.get_all_nslc(datetime=t)
            self.assertEqual(expected_nslc, got_nslc)
            # other dates that have no data should return empty list
            got_nslc = client.get_all_nslc(datetime=t + 2 * 24 * 3600)
            self.assertEqual([], got_nslc)
            got_nslc = client.get_all_nslc(datetime=t - 2 * 24 * 3600)
            self.assertEqual([], got_nslc)


def suite():
    return unittest.makeSuite(SDSTestCase, 'test')


if __name__ == '__main__':
    unittest.main(defaultTest='suite')