diff --git a/py/desitransfer/spacewatch.py b/py/desitransfer/spacewatch.py index d4c6c9a..8c829db 100644 --- a/py/desitransfer/spacewatch.py +++ b/py/desitransfer/spacewatch.py @@ -27,7 +27,12 @@ import pytz utc = pytz.UTC import requests +from desiutil.log import get_logger, DEBUG from . import __version__ as dtVersion +from .common import yesterday + + +log = None class SpacewatchHTMLParser(HTMLParser): @@ -69,7 +74,7 @@ def jpg_list(index): return [index + j for j in parser.jpg_files] -def download_jpg(files, destination): +def download_jpg(files, destination, overwrite=False, test=False): """Download `files` to `destination`. Parameters @@ -78,6 +83,10 @@ def download_jpg(files, destination): A list of URLs to download. destination : :class:`str` A local directory to hold the files. + overwrite : :class:`str`, optional + If ``True``, overwrite any existing files. + test : :class:`bool`, optional + If ``True``, do not download any files. Returns ------- @@ -88,8 +97,9 @@ def download_jpg(files, destination): for jpg in files: base_jpg = jpg.split('/')[-1] dst_jpg = os.path.join(destination, base_jpg) - if os.path.exists(dst_jpg): + if os.path.exists(dst_jpg) and not overwrite: # Overwrite? + log.debug("Skipping existing file: %s.", dst_jpg) pass else: r = requests.get(jpg) @@ -114,6 +124,12 @@ def _options(): prsr = ArgumentParser(description=desc) prsr.add_argument('-d', '--debug', action='store_true', help='Set log level to DEBUG.') + prsr.add_argument('-D', '--date', action='store', metavar='YYYY/MM/DD', + help='Download files for a specific date instead of today.') + prsr.add_argument('-o', '--overwrite', action='store_true', + help='Overwrite any existing files.') + prsr.add_argument('-t', '--test', action='store_true', + help='Do not actually download any files; implies --debug.') prsr.add_argument('-V', '--version', action='version', version='%(prog)s {0}'.format(dtVersion)) prsr.add_argument('destination', metavar='DIR', help='Download files to DIR.') @@ -128,8 +144,26 @@ def main(): :class:`int` An integer suitable for passing to :func:`sys.exit`. """ + global log options = _options() + if options.debug or options.test: + log = get_logger(DEBUG) + else: + log = get_logger() spacewatch_root = 'https://varuna.kpno.noirlab.edu/allsky-all/images/cropped/' - spacewatch_today = spacewatch_root + datetime.date.today().strftime("%Y/%m/%d/") - n_files = download_jpg(jpg_list(spacewatch_today), options.destination) + if options.date is not None: + today = options.date + else: + today = datetime.date.today().strftime("%Y/%m/%d") + y = yesterday() + ystrdy = f"{y[0:4]}/{y[4:6]}/{y[6:8]}" + spacewatch_today = spacewatch_root + today + '/' + spacewatch_yesterday = spacewatch_root + ystrdy + '/' + n_files = download_jpg(jpg_list(spacewatch_today), options.destination, + overwrite=options.overwrite, test=options.test) + log.info("%d files downloaded for %s.", n_files, today) + if options.date is None: + n_files = download_jpg(jpg_list(spacewatch_yesterday), options.destination, + overwrite=options.overwrite, test=options.test) + log.info("%d files downloaded for %s.", n_files, ystrdy) return 0 diff --git a/py/desitransfer/test/test_spacewatch.py b/py/desitransfer/test/test_spacewatch.py index bdb85de..6dfda45 100644 --- a/py/desitransfer/test/test_spacewatch.py +++ b/py/desitransfer/test/test_spacewatch.py @@ -148,21 +148,35 @@ def test_jpg_files(self, mock_requests): 'http://foo.bar/20231031_000405.jpg', 'http://foo.bar/20231031_000605.jpg']) - # @patch('desitransfer.nightwatch.SMTPHandler') - # @patch('desitransfer.nightwatch.RotatingFileHandler') - # @patch('desitransfer.nightwatch.get_logger') - # @patch('desitransfer.nightwatch.log') # Needed to restore the module-level log object after test. - # def test_configure_log(self, mock_log, gl, rfh, smtp): - # """Test logging configuration. - # """ - # with patch.dict('os.environ', - # {'SCRATCH': self.tmp.name, - # 'DESI_ROOT': '/desi/root', - # 'DESI_SPECTRO_DATA': '/desi/root/spectro/data'}): - # with patch.object(sys, 'argv', ['desi_nightwatch_transfer', '--debug']): - # options = _options() - # _configure_log(options) - # rfh.assert_called_once_with('/desi/root/spectro/nightwatch/desi_nightwatch_transfer.log', - # backupCount=100, maxBytes=100000000) - # gl.assert_called_once_with(timestamp=True) - # gl().setLevel.assert_called_once_with(logging.DEBUG) + @patch('desitransfer.spacewatch.log') + @patch('os.utime') + @patch('desitransfer.spacewatch.requests') + @patch('os.path.exists') + def test_download_jpg(self, mock_exists, mock_requests, mock_utime, mock_log): + """Test downloads of JPEG files. + """ + mock_exists.side_effect = [True, False, False, False] + mock_contents = Mock() + mock_contents.headers = {'Last-Modified': 'Mon, 30 Oct 2023 00:00:24 GMT'} + mock_contents.status_code = 200 + mock_contents.content = b"""123456789""" + mock_requests.get.return_value = mock_contents + files = ['http://foo.bar/20231031_000005.jpg', + 'http://foo.bar/20231031_000205.jpg', + 'http://foo.bar/20231031_000405.jpg', + 'http://foo.bar/20231031_000605.jpg'] + destination = self.tmp.name + n = download_jpg(files, destination) + self.assertEqual(n, 3) + mock_exists.assert_has_calls([call(os.path.join(destination, '20231031_000005.jpg')), + call(os.path.join(destination, '20231031_000205.jpg')), + call(os.path.join(destination, '20231031_000405.jpg')), + call(os.path.join(destination, '20231031_000605.jpg'))]) + mock_requests.get.assert_has_calls([call('http://foo.bar/20231031_000205.jpg'), + call('http://foo.bar/20231031_000405.jpg'), + call('http://foo.bar/20231031_000605.jpg')]) + mock_utime.assert_has_calls([call(os.path.join(destination, '20231031_000205.jpg'), (1698624024, 1698624024)), + call(os.path.join(destination, '20231031_000405.jpg'), (1698624024, 1698624024)), + call(os.path.join(destination, '20231031_000605.jpg'), (1698624024, 1698624024))]) + mock_log.debug.assert_has_calls([call("Skipping existing file: %s.", + os.path.join(destination, '20231031_000005.jpg'))])