Merge pull request #149 from alurban/unit-tests

Added unit tests for hveto.cli.trace

hveto/cli/tests/__init__.py

@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) Alex Urban (2020)
+#
+# This file is part of the hveto python package.
+#
+# hveto is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# hveto is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with hveto.  If not, see <http://www.gnu.org/licenses/>.
+
+"""Unit tests for hveto.cli
+"""
+
+__author__ = "Alex Urban <[email protected]>"

hveto/cli/tests/test_trace.py

@@ -0,0 +1,105 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) Alex Urban (2020)
+#
+# This file is part of the hveto python package.
+#
+# hveto is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# hveto is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with hveto.  If not, see <http://www.gnu.org/licenses/>.
+
+"""Tests for :mod:`hveto.cli.trace`
+"""
+
+import os.path
+import pytest
+import shutil
+
+from gwpy.segments import (
+    Segment,
+    SegmentList,
+)
+
+from .. import trace
+
+__author__ = "Alex Urban <[email protected]>"
+
+
+# -- test configuration
+
+VALID_RESULTS = """{
+  "rounds": [{
+    "files": {"VETO_SEGS": ["TEST-HVETO-SEGMENTS.txt"]},
+    "name": "L1:ASC-X_TR_A_NSUM_OUT_DQ",
+    "significance": 78.93832125380759,
+    "snr": 11.0,
+    "window": 0.4
+  }]
+}"""
+
+EMPTY_RESULTS = '{"rounds": [{"files": {"VETO_SEGS": []}}]}'
+TEST_SEGMENTS = SegmentList([Segment(0, 8)])
+
+
+# -- cli tests ----------------------------------------------------------------
+
+def test_main(caplog, tmpdir):
+    indir = str(tmpdir)
+    summary_stats = os.path.join(indir, 'summary-stats.json')
+    print(VALID_RESULTS, file=open(summary_stats, 'w'))
+    TEST_SEGMENTS.write(os.path.join(indir, 'TEST-HVETO-SEGMENTS.txt'))
+    args = [
+        '--trigger-time', '4',
+        '--directory', indir,
+        '--verbose',
+    ]
+    # test output
+    trace.main(args)
+    assert "Running in verbose mode" in caplog.text
+    assert "Search directory: {}".format(indir) in caplog.text
+    assert ("Trigger time 4.0 was vetoed in round 1 "
+            "by segment [0 ... 8)" in caplog.text)
+    assert "Round winner: L1:ASC-X_TR_A_NSUM_OUT_DQ" in caplog.text
+    assert "Significance: 78.93832125380759" in caplog.text
+    assert "SNR: 11.0" in caplog.text
+    assert "Window: 0.4" in caplog.text
+    # clean up
+    shutil.rmtree(indir, ignore_errors=True)
+
+
+def test_main_no_input(caplog, tmpdir):
+    indir = str(tmpdir)
+    args = [
+        '--trigger-time', '0',
+        '--directory', indir,
+    ]
+    # test output
+    with pytest.raises(IOError):
+        trace.main(args)
+    assert ("'summary-stats.json' was not found in the "
+            "input directory" in caplog.text)
+    # clean up
+    shutil.rmtree(indir, ignore_errors=True)
+
+
+def test_main_empty_input(caplog, tmpdir):
+    indir = str(tmpdir)
+    summary_stats = os.path.join(indir, 'summary-stats.json')
+    print(EMPTY_RESULTS, file=open(summary_stats, 'w'))
+    args = [
+        '--trigger-time', '0',
+        '--directory', indir,
+    ]
+    # test output
+    trace.main(args)
+    assert "Trigger time 0.0 was not vetoed" in caplog.text
+    # clean up
+    shutil.rmtree(indir, ignore_errors=True)

hveto/cli/trace.py

@@ -16,7 +16,7 @@
 # You should have received a copy of the GNU General Public License
 # along with hveto.  If not, see <http://www.gnu.org/licenses/>.
 
-"""Check if a trigger was vetoed by a specified hveto run
+"""Check whether a specified trigger time was vetoed by an hveto analysis
 """
 
 import os
@@ -60,7 +60,7 @@ def create_parser():
         '--directory',
         required=True,
         type=_abs_path,
-        help=('path to hveto-generated folder containing '
+        help=('path to the hveto output directory containing '
               'a summary-stats.json file'),
     )
     parser.add_argument(
@@ -78,45 +78,51 @@ def create_parser():
 # -- main code block ----------------------------------------------------------
 
 def main(args=None):
-    """Run the trace tool
+    """Run the hveto-trace command-line tool
     """
     parser = create_parser()
     args = parser.parse_args(args=args)
-    directory = args.directory
 
+    # initialize variables
+    time = float(args.trigger_time)
+    segment = None
+
+    # initialize logger
     logger = cli.logger(name=PROG.split('python -m ').pop(),
                         level=args.loglevel)
     logger.debug('Running in verbose mode')
-    logger.debug('Search directory: %s' % directory)
-
-    trigger_time = float(args.trigger_time)
-    if directory[-1] != '/':
-        directory += '/'
+    logger.debug('Search directory: {}'.format(args.directory))
 
-    try:
-        segment_stats = json.load(open('%ssummary-stats.json' % directory))
+    try:  # read veto segment statistics
+        segment_stats = json.load(open(os.path.join(
+            args.directory, 'summary-stats.json')))
     except IOError:
-        logger.error("'summary-stats.json' was not found "
-                     "in the input directory")
-        sys.exit(0)
+        logger.critical("'summary-stats.json' was not found "
+                        "in the input directory")
+        raise
 
-    for i, cround in enumerate(segment_stats['rounds']):
+    # loop over and log results to output
+    for (i, cround) in enumerate(segment_stats['rounds']):
         seg_files = filter(
-            lambda f_name: '.txt' in f_name, cround[u'files'][u'VETO_SEGS'])
+            lambda f_name: '.txt' in f_name,
+            cround[u'files'][u'VETO_SEGS'])
         for f in seg_files:
-            segments = SegmentList.read(os.path.join(directory, f))
-            for segment in segments:
-                if segment[0] <= trigger_time <= segment[1]:
-                    logger.info('Signal was vetoed in round %d by '
-                                'segment %s' % ((i + 1), segment))
-                    logger.debug('Winner: %s' % cround['name'])
-                    logger.debug('Significance: %s' % cround['significance'])
-                    logger.debug('SNR: %s' % cround['snr'])
-                    logger.debug('Window: %s' % cround['window'])
-                    sys.exit(0)
-
-    logger.info('Signal was not vetoed.')
-
+            segments = SegmentList.read(os.path.join(args.directory, f))
+            if time in segments:
+                segment = segments[segments.find(time)]
+                logger.info('Trigger time {0} was vetoed in round {1} by '
+                            'segment {2}'.format(time, (i + 1), segment))
+                logger.debug('Round winner: {}'.format(cround['name']))
+                logger.debug('Significance: {}'.format(cround['significance']))
+                logger.debug('SNR: {}'.format(cround['snr']))
+                logger.debug('Window: {}'.format(cround['window']))
+
+    if segment is None:
+        # if we got here, the signal was not vetoed
+        logger.info('Trigger time {} was not vetoed'.format(time))
+
+
+# -- run from command-line ----------------------------------------------------
 
 if __name__ == "__main__":
     main()