Coverage for tests/test_packageAlerts.py: 97%

1# This file is part of ap_association. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22import io 

23import os 

24 

25import numpy as np 

26import pandas as pd 

27import tempfile 

28import unittest 

29from unittest.mock import patch, Mock 

30from astropy import wcs 

31from astropy.nddata import CCDData 

32import fastavro 

33try: 

34 import confluent_kafka 

35 from confluent_kafka import KafkaException 

36except ImportError: 

37 confluent_kafka = None 

38 

39import lsst.alert.packet as alertPack 

40from lsst.ap.association import PackageAlertsConfig, PackageAlertsTask 

41from lsst.ap.association.utils import readSchemaFromApdb 

42from lsst.afw.cameraGeom.testUtils import DetectorWrapper 

43import lsst.afw.image as afwImage 

44from lsst.daf.base import DateTime 

45from lsst.dax.apdb import Apdb, ApdbSql 

46import lsst.geom as geom 

47import lsst.meas.base.tests 

48from lsst.sphgeom import Box 

49import lsst.utils.tests 

50from lsst.pipe.tasks.functors import LocalWcs 

51from lsst.pipe.tasks.schemaUtils import convertDataFrameToSdmSchema 

52import utils_tests 

53 

54 

55def _roundTripThroughApdb(objects, sources, forcedSources, dateTime): 

56 """Run object and source catalogs through the Apdb to get the correct 

57 table schemas. 

58 

59 Parameters 

60 ---------- 

61 objects : `pandas.DataFrame` 

62 Set of test DiaObjects to round trip. 

63 sources : `pandas.DataFrame` 

64 Set of test DiaSources to round trip. 

65 forcedSources : `pandas.DataFrame` 

66 Set of test DiaForcedSources to round trip. 

67 dateTime : `astropy.time.Time` 

68 Time for the Apdb. 

69 

70 Returns 

71 ------- 

72 objects : `pandas.DataFrame` 

73 Round tripped objects. 

74 sources : `pandas.DataFrame` 

75 Round tripped sources. 

76 """ 

77 with tempfile.NamedTemporaryFile() as tmpFile: 

78 apdbConfig = ApdbSql.init_database(db_url="sqlite:///" + tmpFile.name) 

79 apdb = Apdb.from_config(apdbConfig) 

80 

81 wholeSky = Box.full() 

82 loadedObjects = apdb.getDiaObjects(wholeSky) 

83 if loadedObjects.empty: 83 ↛ 86line 83 didn't jump to line 86 because the condition on line 83 was always true

84 diaObjects = objects 

85 else: 

86 diaObjects = pd.concat([loadedObjects, objects]) 

87 loadedDiaSources = apdb.getDiaSources(wholeSky, [], dateTime) 

88 if loadedDiaSources.empty: 88 ↛ 91line 88 didn't jump to line 91 because the condition on line 88 was always true

89 diaSources = sources 

90 else: 

91 diaSources = pd.concat([loadedDiaSources, sources]) 

92 loadedDiaForcedSources = apdb.getDiaForcedSources(wholeSky, [], dateTime) 

93 if loadedDiaForcedSources.empty: 93 ↛ 96line 93 didn't jump to line 96 because the condition on line 93 was always true

94 diaForcedSources = forcedSources 

95 else: 

96 diaForcedSources = pd.concat([loadedDiaForcedSources, forcedSources]) 

97 

98 apdb.store(dateTime, diaObjects, diaSources, diaForcedSources) 

99 

100 diaObjects = apdb.getDiaObjects(wholeSky) 

101 diaSources = apdb.getDiaSources(wholeSky, 

102 np.unique(diaObjects["diaObjectId"]), 

103 dateTime) 

104 diaForcedSources = apdb.getDiaForcedSources( 

105 wholeSky, np.unique(diaObjects["diaObjectId"]), dateTime) 

106 

107 diaObjects.set_index("diaObjectId", drop=False, inplace=True) 

108 diaSources.set_index(["diaObjectId", "band", "diaSourceId"], 

109 drop=False, 

110 inplace=True) 

111 diaForcedSources.set_index(["diaObjectId"], drop=False, inplace=True) 

112 

113 # apply SDM type standardization to catch pandas typing issues 

114 schema = readSchemaFromApdb(apdb) 

115 diaObjects = convertDataFrameToSdmSchema(schema, diaObjects, tableName="DiaObject", skipIndex=True) 

116 diaSources = convertDataFrameToSdmSchema(schema, diaSources, tableName="DiaSource", skipIndex=True) 

117 diaForcedSources = convertDataFrameToSdmSchema(schema, diaForcedSources, tableName="DiaForcedSource", 

118 skipIndex=True) 

119 

120 return (diaObjects, diaSources, diaForcedSources) 

121 

122 

123VISIT = 2 

124DETECTOR = 42 

125 

126 

127def mock_alert(dia_source_id): 

128 """Generate a minimal mock alert. 

129 """ 

130 return { 

131 "diaSourceId": dia_source_id, 

132 "diaSource": { 

133 "midpointMjdTai": 5, 

134 "diaSourceId": 1234, 

135 "visit": VISIT, 

136 "detector": DETECTOR, 

137 "band": 'g', 

138 "ra": 12.5, 

139 "dec": -16.9, 

140 # These types are 32-bit floats in the avro schema, so we have to 

141 # make them that type here, so that they round trip appropriately. 

142 "x": np.float32(15.7), 

143 "y": np.float32(89.8), 

144 "apFlux": np.float32(54.85), 

145 "apFluxErr": np.float32(70.0), 

146 "snr": np.float32(6.7), 

147 "psfFlux": np.float32(700.0), 

148 "psfFluxErr": np.float32(90.0), 

149 "psfNdata": int(25), 

150 "dipoleNdata": int(43), 

151 "trailNdata": int(5), 

152 "bboxSize": int(50), 

153 # unlike in transformDiaSourceCatalog.py we need a timezone-aware 

154 # version because mock_alert does not go through pandas 

155 "timeProcessedMjdTai": DateTime.now().get(system=DateTime.MJD, scale=DateTime.TAI) 

156 } 

157 } 

158 

159 

160def mock_ss_alert(dia_source_id, ss_object_id): 

161 """Generate a minimal mock alert. 

162 """ 

163 alert = mock_alert(dia_source_id) 

164 alert['mpc_orbits'] = { 

165 'id': 42, 

166 'designation': '2020 AH11', # a string-typed field 

167 'q': np.float64(0.99999), # a double-typed field 

168 'unpacked_primary_provisional_designation': '2020 AH11', 

169 'packed_primary_provisional_designation': 'J201AH1', 

170 

171 } 

172 alert['ssSource'] = { 

173 'diaSourceId': dia_source_id, 

174 'ssObjectId': ss_object_id, 

175 'eclLambda': np.float64(3.141592), # a double-typed field 

176 'eclBeta': np.float64(3.141592), # a double-typed field 

177 'galLon': np.float64(3.141592), # a double-typed field 

178 'galLat': np.float64(3.141592), # a double-typed field 

179 'helioRange': np.float32(3.141592), # a float-typed field 

180 } 

181 return alert 

182 

183 

184def _deserialize_alert(alert_bytes): 

185 """Deserialize an alert message from Kafka. 

186 

187 Parameters 

188 ---------- 

189 alert_bytes : `bytes` 

190 Binary-encoding serialized Avro alert, including Confluent Wire 

191 Format prefix. 

192 

193 Returns 

194 ------- 

195 alert : `dict` 

196 An alert payload. 

197 """ 

198 schema = alertPack.Schema.from_uri(str(alertPack.get_uri_to_latest_schema())) 

199 content_bytes = io.BytesIO(alert_bytes[5:]) 

200 

201 return fastavro.schemaless_reader(content_bytes, schema.definition) 

202 

203 

204class TestPackageAlerts(lsst.utils.tests.TestCase): 

205 def setUp(self): 

206 # Create an instance of random generator with fixed seed. 

207 rng = np.random.default_rng(1234) 

208 

209 patcher = patch.dict(os.environ, {"AP_KAFKA_PRODUCER_PASSWORD": "fake_password", 

210 "AP_KAFKA_PRODUCER_USERNAME": "fake_username", 

211 "AP_KAFKA_SERVER": "fake_server", 

212 "AP_KAFKA_TOPIC": "fake_topic"}) 

213 self.environ = patcher.start() 

214 self.addCleanup(patcher.stop) 

215 self.cutoutSize = 35 

216 self.center = lsst.geom.Point2D(50.1, 49.8) 

217 self.bbox = lsst.geom.Box2I(lsst.geom.Point2I(-20, -30), 

218 lsst.geom.Extent2I(140, 160)) 

219 self.dataset = lsst.meas.base.tests.TestDataset(self.bbox) 

220 self.dataset.addSource(100000.0, self.center) 

221 exposure, catalog = self.dataset.realize( 

222 10.0, 

223 self.dataset.makeMinimalSchema(), 

224 randomSeed=1234) 

225 self.exposure = exposure 

226 detector = DetectorWrapper(id=DETECTOR, bbox=exposure.getBBox()).detector 

227 self.exposure.setDetector(detector) 

228 

229 visit = afwImage.VisitInfo( 

230 id=VISIT, 

231 exposureTime=200., 

232 date=DateTime("2014-05-13T17:00:00.000000000", 

233 DateTime.Timescale.TAI), 

234 boresightRotAngle=geom.Angle(0.785398)) 

235 self.exposure.info.id = 1234 

236 self.exposure.info.setVisitInfo(visit) 

237 

238 self.exposure.setFilter( 

239 afwImage.FilterLabel(band='g', physical="g.MP9401")) 

240 

241 diaObjects = utils_tests.makeDiaObjects(2, self.exposure, rng) 

242 diaSourceHistory = utils_tests.makeDiaSources( 

243 10, diaObjects["diaObjectId"].to_numpy(), self.exposure, rng) 

244 diaForcedSources = utils_tests.makeDiaForcedSources( 

245 10, diaObjects["diaObjectId"].to_numpy(), self.exposure, rng) 

246 self.diaObjects, diaSourceHistory, self.diaForcedSources = _roundTripThroughApdb( 

247 diaObjects, 

248 diaSourceHistory, 

249 diaForcedSources, 

250 self.exposure.visitInfo.date.toAstropy()) 

251 diaSourceHistory["programId"] = 0 

252 

253 self.diaSources = diaSourceHistory.loc[[(1, "g", 9), (2, "g", 10)], :] 

254 self.diaSources["bboxSize"] = self.cutoutSize 

255 self.diaSourceHistory = diaSourceHistory.drop(labels=[(1, "g", 9), 

256 (2, "g", 10)]) 

257 

258 self.cutoutWcs = wcs.WCS(naxis=2) 

259 self.cutoutWcs.wcs.crpix = [self.center[0], self.center[1]] 

260 self.cutoutWcs.wcs.crval = [ 

261 self.exposure.getWcs().getSkyOrigin().getRa().asDegrees(), 

262 self.exposure.getWcs().getSkyOrigin().getDec().asDegrees()] 

263 self.cutoutWcs.wcs.cd = self.exposure.getWcs().getCdMatrix() 

264 self.cutoutWcs.wcs.ctype = ["RA---TAN", "DEC--TAN"] 

265 

266 # create a mask with planes=("BAD", "SATURATED", "EDGE") 

267 bad_bit = self.exposure.mask.getMaskPlane("BAD") 

268 self.exposure.mask.array[50:55, 100:110] |= bad_bit 

269 saturated_bit = self.exposure.mask.getMaskPlane("SAT") 

270 self.exposure.mask.array[10:25, 60:68] |= saturated_bit 

271 edge_bit = self.exposure.mask.getMaskPlane("EDGE") 

272 self.exposure.mask.array[0:5, ...] |= edge_bit 

273 

274 def testCreateExtentMinimum(self): 

275 """Test the extent creation for the cutout bbox returns a cutout with 

276 the minimum cutouut size. 

277 """ 

278 packConfig = PackageAlertsConfig() 

279 # Just create a minimum less than the default cutout size. 

280 packConfig.minCutoutSize = self.cutoutSize - 5 

281 packageAlerts = PackageAlertsTask(config=packConfig) 

282 extent = packageAlerts.createDiaSourceExtent( 

283 packConfig.minCutoutSize - 5) 

284 self.assertTrue(extent == geom.Extent2I(packConfig.minCutoutSize, 

285 packConfig.minCutoutSize)) 

286 # Test that the cutout size is correctly increased. 

287 extent = packageAlerts.createDiaSourceExtent(self.cutoutSize) 

288 self.assertTrue(extent == geom.Extent2I(self.cutoutSize, 

289 self.cutoutSize)) 

290 

291 def testCreateExtentMaximum(self): 

292 """Test the extent creation for the cutout bbox returns a cutout with 

293 the maximum cutout size. 

294 """ 

295 packConfig = PackageAlertsConfig() 

296 # Just create a maximum more than the default cutout size. 

297 packConfig.maxCutoutSize = self.cutoutSize + 5 

298 packageAlerts = PackageAlertsTask(config=packConfig) 

299 extent = packageAlerts.createDiaSourceExtent( 

300 packConfig.maxCutoutSize + 5) 

301 self.assertTrue(extent == geom.Extent2I(packConfig.maxCutoutSize, 

302 packConfig.maxCutoutSize)) 

303 # Test that the cutout size is correctly reduced. 

304 extent = packageAlerts.createDiaSourceExtent(self.cutoutSize) 

305 self.assertTrue(extent == geom.Extent2I(self.cutoutSize, 

306 self.cutoutSize)) 

307 

308 def testPsfSerializedAsFloat32Bitpix(self): 

309 """Test that PSF is downcast to float32 and serialized with BITPIX=-32.""" 

310 from astropy.io import fits 

311 

312 packageAlerts = PackageAlertsTask() 

313 

314 diaSrcId = 1234 

315 ccdData = packageAlerts.createCcdDataCutout( 

316 self.exposure, 

317 self.exposure.getWcs().getSkyOrigin(), 

318 self.exposure.getWcs().getPixelOrigin(), 

319 self.exposure.getBBox().getDimensions(), 

320 self.exposure.getPhotoCalib(), 

321 diaSrcId) 

322 

323 self.assertIsNotNone(ccdData) 

324 self.assertIsNotNone(ccdData.psf) 

325 self.assertEqual(ccdData.psf.dtype, np.float32) 

326 

327 cutoutBytes = packageAlerts.streamCcdDataToBytes(ccdData) 

328 

329 with io.BytesIO(cutoutBytes) as bytesIO: 

330 with fits.open(bytesIO) as hdul: 

331 self.assertIn("PSFIMAGE", hdul) 

332 self.assertEqual(hdul["PSFIMAGE"].header["BITPIX"], -32) 

333 

334 with io.BytesIO(cutoutBytes) as bytesIO: 

335 cutoutFromBytes = CCDData.read(bytesIO, format="fits") 

336 self.assertIsNotNone(cutoutFromBytes.psf) 

337 self.assertEqual(cutoutFromBytes.psf.dtype, '>f4') 

338 

339 def testCreateCcdDataCutout(self): 

340 """Test that the data is being extracted into the CCDData cutout 

341 correctly. 

342 """ 

343 packageAlerts = PackageAlertsTask() 

344 

345 diaSrcId = 1234 

346 ccdData = packageAlerts.createCcdDataCutout( 

347 self.exposure, 

348 self.exposure.getWcs().getSkyOrigin(), 

349 self.exposure.getWcs().getPixelOrigin(), 

350 self.exposure.getBBox().getDimensions(), 

351 self.exposure.getPhotoCalib(), 

352 diaSrcId) 

353 calibExposure = self.exposure.getPhotoCalib().calibrateImage( 

354 self.exposure.getMaskedImage()) 

355 

356 self.assertFloatsAlmostEqual(ccdData.wcs.wcs.cd, 

357 self.cutoutWcs.wcs.cd) 

358 self.assertFloatsAlmostEqual(ccdData.data, 

359 calibExposure.getImage().array) 

360 self.assertFloatsAlmostEqual( 

361 ccdData.psf, 

362 self.exposure.psf.computeKernelImage(self.center).array.astype(np.float32), 

363 rtol=1e-6, atol=1e-6 

364 ) 

365 self.assertFloatsAlmostEqual(self.exposure.mask.array, ccdData.flags) 

366 

367 cutoutBytes = packageAlerts.streamCcdDataToBytes(ccdData) 

368 readCcdData = CCDData.read(io.BytesIO(cutoutBytes), format="fits", hdu_flags='MASKPLANE') 

369 self.assertFloatsAlmostEqual(self.exposure.mask.array, readCcdData.flags) 

370 self.assertFloatsAlmostEqual(readCcdData.data, ccdData.data) 

371 self.assertFloatsAlmostEqual(readCcdData.flags, ccdData.flags) 

372 

373 ccdData = packageAlerts.createCcdDataCutout( 

374 self.exposure, 

375 geom.SpherePoint(0, 0, geom.degrees), 

376 geom.Point2D(-100, -100), 

377 self.exposure.getBBox().getDimensions(), 

378 self.exposure.getPhotoCalib(), 

379 diaSrcId) 

380 self.assertTrue(ccdData is None) 

381 

382 def testMakeLocalTransformMatrix(self): 

383 """Test that the local WCS approximation is correct. 

384 """ 

385 packageAlerts = PackageAlertsTask() 

386 

387 sphPoint = self.exposure.getWcs().pixelToSky(self.center) 

388 cutout = self.exposure.getCutout(sphPoint, 

389 geom.Extent2I(self.cutoutSize, 

390 self.cutoutSize)) 

391 cd = packageAlerts.makeLocalTransformMatrix( 

392 cutout.getWcs(), self.center, sphPoint) 

393 self.assertFloatsAlmostEqual( 

394 cd, 

395 cutout.getWcs().getCdMatrix(), 

396 rtol=5e-10, 

397 atol=5e-10) 

398 

399 def testStreamCcdDataToBytes(self): 

400 """Test round tripping an CCDData cutout to bytes and back. 

401 """ 

402 packageAlerts = PackageAlertsTask() 

403 

404 sphPoint = self.exposure.getWcs().pixelToSky(self.center) 

405 cutout = self.exposure.getCutout(sphPoint, 

406 geom.Extent2I(self.cutoutSize, 

407 self.cutoutSize)) 

408 cutoutCcdData = CCDData( 

409 data=cutout.getImage().array, 

410 wcs=self.cutoutWcs, 

411 unit="adu") 

412 

413 cutoutBytes = packageAlerts.streamCcdDataToBytes(cutoutCcdData) 

414 with io.BytesIO(cutoutBytes) as bytesIO: 

415 cutoutFromBytes = CCDData.read(bytesIO, format="fits") 

416 self.assertFloatsAlmostEqual(cutoutCcdData.data, cutoutFromBytes.data) 

417 

418 def testMakeAlertDict(self): 

419 """Test stripping data from the various data products and into a 

420 dictionary "alert". 

421 """ 

422 packageAlerts = PackageAlertsTask() 

423 dia_source_id = 1234 

424 

425 for srcIdx, diaSource in self.diaSources.iterrows(): 

426 sphPoint = geom.SpherePoint(diaSource["ra"], 

427 diaSource["dec"], 

428 geom.degrees) 

429 pixelPoint = geom.Point2D(diaSource["x"], diaSource["y"]) 

430 cutout = self.exposure.getCutout(sphPoint, 

431 geom.Extent2I(self.cutoutSize, 

432 self.cutoutSize)) 

433 ccdCutout = packageAlerts.createCcdDataCutout( 

434 cutout, 

435 sphPoint, 

436 pixelPoint, 

437 geom.Extent2I(self.cutoutSize, self.cutoutSize), 

438 cutout.getPhotoCalib(), 

439 1234, 

440 rotPa=cutout.visitInfo.boresightRotAngle.asDegrees()) 

441 cutoutBytes = packageAlerts.streamCcdDataToBytes( 

442 ccdCutout) 

443 objSources = self.diaSourceHistory.loc[srcIdx[0]] 

444 objForcedSources = self.diaForcedSources.loc[srcIdx[0]] 

445 alert = packageAlerts.makeAlertDict( 

446 dia_source_id, 

447 self.exposure.visitInfo.getObservationReason(), 

448 self.exposure.visitInfo.getObject(), 

449 diaSource, 

450 self.diaObjects.loc[srcIdx[0]], 

451 objSources, 

452 objForcedSources, 

453 ccdCutout, 

454 ccdCutout, 

455 ccdCutout) 

456 self.assertEqual(len(alert), 13) 

457 

458 self.assertEqual(alert["diaSourceId"], dia_source_id) 

459 self.assertEqual(alert["diaSource"], diaSource.to_dict()) 

460 self.assertIsNone(alert["observation_reason"]) 

461 self.assertIsNone(alert["target_name"]) 

462 self.assertEqual(alert["cutoutDifference"], 

463 cutoutBytes) 

464 self.assertEqual(alert["cutoutScience"], 

465 cutoutBytes) 

466 self.assertEqual(alert["cutoutTemplate"], 

467 cutoutBytes) 

468 science_cutout = CCDData.read(io.BytesIO(alert["cutoutScience"]), 

469 format="fits") 

470 template_cutout = CCDData.read(io.BytesIO(alert["cutoutTemplate"]), 

471 format="fits") 

472 self.assertAlmostEqual(science_cutout.header["ROTPA"], 

473 template_cutout.header["ROTPA"]) 

474 

475 def testMakeAlertDictSchedulerFields(self): 

476 """Test non-null scheduler fields pass through as expected. 

477 

478 """ 

479 packageAlerts = PackageAlertsTask() 

480 dia_source_id = 1234 

481 

482 for srcIdx, diaSource in self.diaSources.iterrows(): 

483 objSources = self.diaSourceHistory.loc[srcIdx[0]] 

484 objForcedSources = self.diaForcedSources.loc[srcIdx[0]] 

485 obs_reason = f"obs_reason_{srcIdx}", 

486 target = f"target_name_{srcIdx}", 

487 alert = packageAlerts.makeAlertDict( 

488 dia_source_id, 

489 obs_reason, 

490 target, 

491 diaSource, 

492 self.diaObjects.loc[srcIdx[0]], 

493 objSources, 

494 objForcedSources, 

495 None, 

496 None, 

497 None) 

498 self.assertEqual(len(alert), 13) 

499 

500 self.assertEqual(alert["observation_reason"], obs_reason) 

501 self.assertEqual(alert["target_name"], target) 

502 

503 def testCutoutRotpa(self): 

504 """Test that the ROTPA header keyword matches the boresightRotAngle from visitInfo. 

505 """ 

506 packageAlerts = PackageAlertsTask() 

507 

508 # Create a cutout using existing test exposure 

509 sphPoint = self.exposure.getWcs().pixelToSky(self.center) 

510 cutout = self.exposure.getCutout(sphPoint, 

511 geom.Extent2I(self.cutoutSize, 

512 self.cutoutSize)) 

513 

514 # Create CCDData cutout 

515 ccdCutout = packageAlerts.createCcdDataCutout( 

516 cutout, 

517 sphPoint, 

518 self.center, 

519 geom.Extent2I(self.cutoutSize, self.cutoutSize), 

520 cutout.getPhotoCalib(), 

521 1234, 

522 rotPa=cutout.visitInfo.boresightRotAngle.asDegrees()) 

523 

524 wcs = self.exposure.wcs 

525 cd_matrix = wcs.getCdMatrix() # This gets the CD matrix elements 

526 angle_rad = LocalWcs.computePositionAngle([0], 

527 [cd_matrix[0, 0]], # CD1_1 

528 [cd_matrix[0, 1]], # CD1_2 

529 [cd_matrix[1, 0]], # CD2_1 

530 [cd_matrix[1, 1]] # CD2_2 

531 ) 

532 wcs_pa = geom.Angle(angle_rad.iloc[0], geom.radians).asDegrees() 

533 

534 cutoutBytes = packageAlerts.streamCcdDataToBytes(ccdCutout) 

535 with io.BytesIO(cutoutBytes) as bytesIO: 

536 cutoutFromBytes = CCDData.read(bytesIO, format="fits") 

537 

538 self.assertIn('ROTPA', cutoutFromBytes.header) 

539 self.assertAlmostEqual( 

540 cutoutFromBytes.header['ROTPA'], 

541 self.exposure.visitInfo.boresightRotAngle.asDegrees()) 

542 self.assertAlmostEqual( 

543 cutoutFromBytes.header['ROTPA'], 

544 wcs_pa, places=4) 

545 

546 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled') 

547 def test_produceAlerts_empty_password(self): 

548 """ Test that produceAlerts raises if the password is empty or missing. 

549 """ 

550 self.environ['AP_KAFKA_PRODUCER_PASSWORD'] = "" 

551 with self.assertRaisesRegex(ValueError, "Kafka password"): 

552 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

553 PackageAlertsTask(config=packConfig) 

554 

555 del self.environ['AP_KAFKA_PRODUCER_PASSWORD'] 

556 with self.assertRaisesRegex(ValueError, "Kafka password"): 

557 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

558 PackageAlertsTask(config=packConfig) 

559 

560 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled') 

561 def test_produceAlerts_empty_username(self): 

562 """ Test that produceAlerts raises if the username is empty or missing. 

563 """ 

564 self.environ['AP_KAFKA_PRODUCER_USERNAME'] = "" 

565 with self.assertRaisesRegex(ValueError, "Kafka username"): 

566 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

567 PackageAlertsTask(config=packConfig) 

568 

569 del self.environ['AP_KAFKA_PRODUCER_USERNAME'] 

570 with self.assertRaisesRegex(ValueError, "Kafka username"): 

571 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

572 PackageAlertsTask(config=packConfig) 

573 

574 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled') 

575 def test_produceAlerts_empty_server(self): 

576 """ Test that produceAlerts raises if the server is empty or missing. 

577 """ 

578 self.environ['AP_KAFKA_SERVER'] = "" 

579 with self.assertRaisesRegex(ValueError, "Kafka server"): 

580 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

581 PackageAlertsTask(config=packConfig) 

582 

583 del self.environ['AP_KAFKA_SERVER'] 

584 with self.assertRaisesRegex(ValueError, "Kafka server"): 

585 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

586 PackageAlertsTask(config=packConfig) 

587 

588 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled') 

589 def test_produceAlerts_empty_topic(self): 

590 """ Test that produceAlerts raises if the topic is empty or missing. 

591 """ 

592 self.environ['AP_KAFKA_TOPIC'] = "" 

593 with self.assertRaisesRegex(ValueError, "Kafka topic"): 

594 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

595 PackageAlertsTask(config=packConfig) 

596 

597 del self.environ['AP_KAFKA_TOPIC'] 

598 with self.assertRaisesRegex(ValueError, "Kafka topic"): 

599 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

600 PackageAlertsTask(config=packConfig) 

601 

602 @patch('confluent_kafka.Producer') 

603 @patch.object(PackageAlertsTask, '_server_check') 

604 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled') 

605 def test_produceAlerts_success(self, mock_server_check, mock_producer): 

606 """ Test that produceAlerts calls the producer on all provided alerts 

607 when the alerts are all under the batch size limit. 

608 """ 

609 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

610 packageAlerts = PackageAlertsTask(config=packConfig) 

611 alerts = [mock_alert(1), mock_alert(2), mock_ss_alert(3, 3)] 

612 

613 # Create a variable and assign it an instance of the patched kafka producer 

614 producer_instance = mock_producer.return_value 

615 producer_instance.produce = Mock() 

616 producer_instance.flush = Mock() 

617 unix_midpoint = self.exposure.visitInfo.date.toAstropy().tai.unix 

618 exposure_time = self.exposure.visitInfo.exposureTime 

619 packageAlerts.produceAlerts(alerts, VISIT, DETECTOR, unix_midpoint, exposure_time) 

620 

621 self.assertEqual(mock_server_check.call_count, 1) 

622 self.assertEqual(producer_instance.produce.call_count, len(alerts)) 

623 self.assertEqual(producer_instance.flush.call_count, len(alerts)+1) 

624 

625 @patch('confluent_kafka.Producer') 

626 @patch.object(PackageAlertsTask, '_server_check') 

627 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled') 

628 def test_produceAlerts_one_failure(self, mock_server_check, mock_producer): 

629 """ Test that produceAlerts correctly fails on one alert 

630 and is writing the failure to disk. 

631 """ 

632 counter = 0 

633 

634 def mock_produce(*args, **kwargs): 

635 nonlocal counter 

636 counter += 1 

637 if counter == 2: 

638 raise KafkaException 

639 else: 

640 return 

641 

642 packConfig = PackageAlertsConfig(doProduceAlerts=True, doWriteFailedAlerts=True) 

643 packageAlerts = PackageAlertsTask(config=packConfig) 

644 

645 patcher = patch("builtins.open") 

646 patch_open = patcher.start() 

647 alerts = [mock_alert(1), mock_alert(2), mock_alert(3), mock_ss_alert(4, 4)] 

648 unix_midpoint = self.exposure.visitInfo.date.toAstropy().tai.unix 

649 exposure_time = self.exposure.visitInfo.exposureTime 

650 

651 producer_instance = mock_producer.return_value 

652 producer_instance.produce = Mock(side_effect=mock_produce) 

653 producer_instance.flush = Mock() 

654 packageAlerts.produceAlerts(alerts, VISIT, DETECTOR, unix_midpoint, exposure_time) 

655 

656 self.assertEqual(mock_server_check.call_count, 1) 

657 self.assertEqual(producer_instance.produce.call_count, len(alerts)) 

658 self.assertEqual(patch_open.call_count, 1) 

659 self.assertIn(f"{VISIT}_{DETECTOR}_2.avro", patch_open.call_args.args[0]) 

660 # Because one produce raises, we call flush one fewer times than in the success 

661 # test above. 

662 self.assertEqual(producer_instance.flush.call_count, len(alerts)) 

663 patcher.stop() 

664 

665 @patch.object(PackageAlertsTask, '_server_check') 

666 def testRun_without_produce(self, mock_server_check): 

667 """Test the run method of package alerts with produce set to False and 

668 doWriteAlerts set to true. 

669 """ 

670 packConfig = PackageAlertsConfig(doWriteAlerts=True) 

671 with tempfile.TemporaryDirectory(prefix='alerts') as tempdir: 

672 packConfig.alertWriteLocation = tempdir 

673 packageAlerts = PackageAlertsTask(config=packConfig) 

674 

675 packageAlerts.run(self.diaSources, 

676 self.diaObjects, 

677 self.diaSourceHistory, 

678 self.diaForcedSources, 

679 self.exposure, 

680 self.exposure, 

681 self.exposure) 

682 

683 self.assertEqual(mock_server_check.call_count, 0) 

684 

685 with open(os.path.join(tempdir, f"{VISIT}_{DETECTOR}.avro"), 'rb') as f: 

686 writer_schema, data_stream = \ 

687 packageAlerts.alertSchema.retrieve_alerts(f) 

688 data = list(data_stream) 

689 

690 self.assertEqual(len(data), len(self.diaSources)) 

691 for idx, alert in enumerate(data): 

692 for key, value in alert["diaSource"].items(): 

693 if isinstance(value, float): 

694 if np.isnan(self.diaSources.iloc[idx][key]): 

695 self.assertTrue(np.isnan(value)) 

696 else: 

697 self.assertAlmostEqual( 

698 1 - value / self.diaSources.iloc[idx][key], 

699 0.) 

700 elif value is None: 

701 self.assertTrue(pd.isna(self.diaSources.iloc[idx][key])) 

702 else: 

703 self.assertEqual(value, self.diaSources.iloc[idx][key]) 

704 sphPoint = geom.SpherePoint(alert["diaSource"]["ra"], 

705 alert["diaSource"]["dec"], 

706 geom.degrees) 

707 pixelPoint = geom.Point2D(alert["diaSource"]["x"], alert["diaSource"]["y"]) 

708 cutout = self.exposure.getCutout(sphPoint, 

709 geom.Extent2I(self.cutoutSize, 

710 self.cutoutSize)) 

711 ccdCutout = packageAlerts.createCcdDataCutout( 

712 cutout, 

713 sphPoint, 

714 pixelPoint, 

715 geom.Extent2I(self.cutoutSize, self.cutoutSize), 

716 cutout.getPhotoCalib(), 

717 1234, 

718 rotPa=cutout.visitInfo.boresightRotAngle.asDegrees()) 

719 

720 self.assertEqual(alert["cutoutDifference"],