Skip to content

Commit

Permalink
additional spacewatch tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@weaverba137
weaverba137 committed Nov 1, 2023
1 parent 0369b26 commit 6f96186
Show file tree
Hide file tree
Showing 2 changed files with 70 additions and 22 deletions.
42 changes: 38 additions & 4 deletions py/desitransfer/spacewatch.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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
-------
Expand All @@ -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)
Expand All @@ -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.')
Expand All @@ -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
50 changes: 32 additions & 18 deletions py/desitransfer/test/test_spacewatch.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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'))])

0 comments on commit 6f96186

Please sign in to comment.