import mock
import pytest
from datetime import datetime, timedelta
from unittest.mock import ANY
from astropy.io.fits import Header
from celery.exceptions import Retry
from banzai.celery import stack_calibrations, schedule_calibration_stacking
from banzai.settings import CALIBRATION_STACK_DELAYS
from banzai.utils import date_utils
from banzai.context import Context
from banzai.tests.utils import FakeInstrument, FakeLCOObservationFrame, FakeCCDData
pytestmark = pytest.mark.celery
# TODO: update tests to use same mock lake data as e2e tests
fake_blocks_response_json = {
"results": [
{
"end": "2019-02-19T21:55:09",
"telescope": "2m0a",
"request": {
"configurations": [
{
"priority": 0,
"instrument_type": "2M0-SCICAM-SPECTRAL",
"instrument_configs": [
{
"exposure_time": 0.01,
"exposure_count": 2,
"rotator_mode": "",
"optical_elements": {},
"mode": "",
"extra_params": {},
"bin_y": 1,
"bin_x": 1
}
],
"target": {
"type": "ICRS",
"name": "my target",
"extra_params": {}
},
"acquisition_config": {
"mode": "OFF",
"extra_params": {}
},
"extra_params": {},
"type": "BIAS",
"guiding_config": {
"optical_elements": {},
"exposure_time": 10,
"optional": True,
"mode": "ON",
"extra_params": {}
},
"constraints": {
"max_airmass": 20,
"extra_params": {},
"min_lunar_distance": 0
}
},
{
"priority": 1,
"instrument_type": "2M0-SCICAM-SPECTRAL",
"instrument_configs": [
{
"exposure_time": 0.01,
"exposure_count": 2,
"rotator_mode": "",
"optical_elements": {},
"mode": "",
"extra_params": {},
"bin_y": 1,
"bin_x": 1
}
],
"target": {
"type": "ICRS",
"name": "my target",
"extra_params": {}
},
"acquisition_config": {
"mode": "OFF",
"extra_params": {}
},
"extra_params": {},
"type": "SKY_FLAT",
"guiding_config": {
"optical_elements": {},
"exposure_time": 10,
"optional": True,
"mode": "ON",
"extra_params": {}
},
"constraints": {
"max_airmass": 20,
"extra_params": {},
"min_lunar_distance": 0
}
}
]
},
"site": "coj",
"start": "2019-02-19T20:27:49",
"state": "PENDING",
"proposal": "calibrate",
"enclosure": "clma",
"name": ""
},
{
"end": "2019-02-20T09:55:09",
"telescope": "2m0a",
"request": {
"configurations": [
{
"priority": 0,
"instrument_type": "2M0-SCICAM-SPECTRAL",
"instrument_configs": [
{
"exposure_time": 0.01,
"exposure_count": 2,
"rotator_mode": "",
"optical_elements": {},
"mode": "",
"extra_params": {},
"bin_y": 1,
"bin_x": 1
}
],
"target": {
"type": "ICRS",
"name": "my target",
"extra_params": {}
},
"acquisition_config": {
"mode": "OFF",
"extra_params": {}
},
"extra_params": {},
"type": "BIAS",
"guiding_config": {
"optical_elements": {},
"exposure_time": 10,
"optional": True,
"mode": "ON",
"extra_params": {}
},
"constraints": {
"max_airmass": 20,
"extra_params": {},
"min_lunar_distance": 0
}
}
]
},
"site": "coj",
"start": "2019-02-20T08:27:49",
"state": "PENDING",
"proposal": "calibrate",
"enclosure": "clma",
"name": ""
}
],
}
fake_instruments_response = FakeInstrument()
class TestMain():
@pytest.fixture(scope='function')
def setup(self):
self.site = 'coj'
self.min_date = '2019-02-19T20:27:49'
self.max_date = '2019-02-20T09:55:09'
self.context = Context({'db_address': 'db_address', 'CALIBRATION_IMAGE_TYPES': ['BIAS'],
'CALIBRATION_STACK_DELAYS': {'BIAS': 300},
'CALIBRATION_STACKER_STAGES': {'BIAS': ['banzai.bias.BiasMaker']}})
self.frame_type = 'BIAS'
self.fake_blocks_response_json = fake_blocks_response_json
self.fake_inst = FakeInstrument(site='coj', camera='2m0-SciCam-Spectral', enclosure='clma', telescope='2m0a')
@mock.patch('banzai.celery.stack_calibrations.apply_async')
@mock.patch('banzai.celery.dbs.get_instruments_at_site')
@mock.patch('banzai.celery.get_calibration_blocks_for_time_range')
@mock.patch('banzai.celery.filter_calibration_blocks_for_type')
def test_submit_stacking_tasks_to_queue_no_delay(self, mock_filter_blocks, mock_get_blocks, mock_get_instruments,
mock_stack_calibrations, setup):
mock_get_instruments.return_value = [self.fake_inst]
mock_get_blocks.return_value = self.fake_blocks_response_json
mock_filter_blocks.return_value = [block for block in self.fake_blocks_response_json['results']]
schedule_calibration_stacking(self.site, self.context, self.min_date, self.max_date)
mock_stack_calibrations.assert_called_with(args=(self.min_date, self.max_date, self.fake_inst.id,
self.frame_type, vars(self.context),
mock_filter_blocks.return_value),
countdown=0)
@mock.patch('banzai.celery.stack_calibrations.apply_async')
@mock.patch('banzai.celery.dbs.get_instruments_at_site')
@mock.patch('banzai.celery.get_calibration_blocks_for_time_range')
@mock.patch('banzai.celery.filter_calibration_blocks_for_type')
def test_submit_stacking_tasks_to_queue_with_delay(self, mock_filter_blocks, mock_get_blocks, mock_get_instruments,
mock_stack_calibrations, setup):
mock_get_instruments.return_value = [self.fake_inst]
self.fake_blocks_response_json['results'][0]['end'] = datetime.strftime(datetime.utcnow() + timedelta(minutes=1),
date_utils.TIMESTAMP_FORMAT)
mock_get_blocks.return_value = self.fake_blocks_response_json
mock_filter_blocks.return_value = [block for block in self.fake_blocks_response_json['results']]
schedule_calibration_stacking(self.site, self.context, self.min_date, self.max_date)
mock_stack_calibrations.assert_called_with(args=(self.min_date, self.max_date, self.fake_inst.id,
self.frame_type, vars(self.context),
mock_filter_blocks.return_value),
countdown=(60+CALIBRATION_STACK_DELAYS['BIAS']))
@mock.patch('banzai.calibrations.make_master_calibrations')
@mock.patch('banzai.celery.dbs.get_individual_cal_frames')
@mock.patch('banzai.celery.dbs.get_instrument_by_id')
def test_stack_calibrations(self, mock_get_instrument, mock_get_calibration_images, mock_make_master_cals, setup):
mock_get_instrument.return_value = self.fake_inst
nx, ny = 102, 105
header = {'DATASEC': f'[1:{nx},1:{ny}]', 'DETSEC': f'[1:{nx},1:{ny}]', 'CCDSUM': '1 1',
'OBSTYPE': 'TEST', 'RDNOISE': 3.0, 'TELESCOP': '1m0-02', 'DAY-OBS': '20191209',
'DATE-OBS': '2019-12-09T00:00:00'}
mock_get_calibration_images.return_value = [FakeLCOObservationFrame(hdu_list=[FakeCCDData(meta=Header(header))])
for i in range(2)]
stack_calibrations(self.min_date, self.max_date, 1, self.frame_type, self.context,
[self.fake_blocks_response_json['results'][0]])
mock_make_master_cals.assert_called_with(self.fake_inst, self.frame_type, self.min_date, self.max_date, ANY)
@mock.patch('banzai.calibrations.make_master_calibrations')
@mock.patch('banzai.celery.dbs.get_individual_cal_frames')
@mock.patch('banzai.celery.dbs.get_instrument_by_id')
def test_stack_calibrations_not_enough_images(self, mock_get_instrument, mock_get_calibration_images, mock_maker, setup):
mock_get_instrument.return_value = self.fake_inst
mock_get_calibration_images.return_value = [FakeLCOObservationFrame(hdu_list=[FakeCCDData()])]
with pytest.raises(Retry) as e:
stack_calibrations(self.min_date, self.max_date, 1, self.frame_type, self.context,
[self.fake_blocks_response_json['results'][0]])
assert e.type is Retry
