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 "templateInstrument": "TESTCam", 

154 "templateVersion": 0, 

155 "mixedInstruments": False, 

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

157 # version because mock_alert does not go through pandas 

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

159 } 

160 } 

161 

162 

163def mock_ss_alert(dia_source_id, ss_object_id): 

164 """Generate a minimal mock alert. 

165 """ 

166 alert = mock_alert(dia_source_id) 

167 alert['mpc_orbits'] = { 

168 'id': 42, 

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

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

171 'unpacked_primary_provisional_designation': '2020 AH11', 

172 'packed_primary_provisional_designation': 'J201AH1', 

173 

174 } 

175 alert['ssSource'] = { 

176 'diaSourceId': dia_source_id, 

177 'ssObjectId': ss_object_id, 

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

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

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

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

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

183 } 

184 return alert 

185 

186 

187def _deserialize_alert(alert_bytes): 

188 """Deserialize an alert message from Kafka. 

189 

190 Parameters 

191 ---------- 

192 alert_bytes : `bytes` 

193 Binary-encoding serialized Avro alert, including Confluent Wire 

194 Format prefix. 

195 

196 Returns 

197 ------- 

198 alert : `dict` 

199 An alert payload. 

200 """ 

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

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

203 

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

205 

206 

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

208 def setUp(self): 

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

210 rng = np.random.default_rng(1234) 

211 

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

213 "AP_KAFKA_PRODUCER_USERNAME": "fake_username", 

214 "AP_KAFKA_SERVER": "fake_server", 

215 "AP_KAFKA_TOPIC": "fake_topic"}) 

216 self.environ = patcher.start() 

217 self.addCleanup(patcher.stop) 

218 self.cutoutSize = 35 

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

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

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

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

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

224 exposure, catalog = self.dataset.realize( 

225 10.0, 

226 self.dataset.makeMinimalSchema(), 

227 randomSeed=1234) 

228 self.exposure = exposure 

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

230 self.exposure.setDetector(detector) 

231 

232 visit = afwImage.VisitInfo( 

233 id=VISIT, 

234 exposureTime=200., 

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

236 DateTime.Timescale.TAI), 

237 boresightRotAngle=geom.Angle(0.785398)) 

238 self.exposure.info.id = 1234 

239 self.exposure.info.setVisitInfo(visit) 

240 

241 self.exposure.setFilter( 

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

243 

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

245 diaSourceHistory = utils_tests.makeDiaSources( 

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

247 diaForcedSources = utils_tests.makeDiaForcedSources( 

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

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

250 diaObjects, 

251 diaSourceHistory, 

252 diaForcedSources, 

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

254 diaSourceHistory["programId"] = 0 

255 

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

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

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

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

260 

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

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

263 self.cutoutWcs.wcs.crval = [ 

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

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

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

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

268 

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

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

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

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

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

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

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

276 

277 def testCreateExtentMinimum(self): 

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

279 the minimum cutouut size. 

280 """ 

281 packConfig = PackageAlertsConfig() 

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

283 packConfig.minCutoutSize = self.cutoutSize - 5 

284 packageAlerts = PackageAlertsTask(config=packConfig) 

285 extent = packageAlerts.createDiaSourceExtent( 

286 packConfig.minCutoutSize - 5) 

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

288 packConfig.minCutoutSize)) 

289 # Test that the cutout size is correctly increased. 

290 extent = packageAlerts.createDiaSourceExtent(self.cutoutSize) 

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

292 self.cutoutSize)) 

293 

294 def testCreateExtentMaximum(self): 

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

296 the maximum cutout size. 

297 """ 

298 packConfig = PackageAlertsConfig() 

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

300 packConfig.maxCutoutSize = self.cutoutSize + 5 

301 packageAlerts = PackageAlertsTask(config=packConfig) 

302 extent = packageAlerts.createDiaSourceExtent( 

303 packConfig.maxCutoutSize + 5) 

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

305 packConfig.maxCutoutSize)) 

306 # Test that the cutout size is correctly reduced. 

307 extent = packageAlerts.createDiaSourceExtent(self.cutoutSize) 

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

309 self.cutoutSize)) 

310 

311 def testPsfSerializedAsFloat32Bitpix(self): 

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

313 from astropy.io import fits 

314 

315 packageAlerts = PackageAlertsTask() 

316 

317 diaSrcId = 1234 

318 ccdData = packageAlerts.createCcdDataCutout( 

319 self.exposure, 

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

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

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

323 self.exposure.getPhotoCalib(), 

324 diaSrcId) 

325 

326 self.assertIsNotNone(ccdData) 

327 self.assertIsNotNone(ccdData.psf) 

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

329 

330 cutoutBytes = packageAlerts.streamCcdDataToBytes(ccdData) 

331 

332 with io.BytesIO(cutoutBytes) as bytesIO: 

333 with fits.open(bytesIO) as hdul: 

334 self.assertIn("PSFIMAGE", hdul) 

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

336 

337 with io.BytesIO(cutoutBytes) as bytesIO: 

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

339 self.assertIsNotNone(cutoutFromBytes.psf) 

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

341 

342 def testCreateCcdDataCutout(self): 

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

344 correctly. 

345 """ 

346 packageAlerts = PackageAlertsTask() 

347 

348 diaSrcId = 1234 

349 ccdData = packageAlerts.createCcdDataCutout( 

350 self.exposure, 

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

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

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

354 self.exposure.getPhotoCalib(), 

355 diaSrcId) 

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

357 self.exposure.getMaskedImage()) 

358 

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

360 self.cutoutWcs.wcs.cd) 

361 self.assertFloatsAlmostEqual(ccdData.data, 

362 calibExposure.getImage().array) 

363 self.assertFloatsAlmostEqual( 

364 ccdData.psf, 

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

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

367 ) 

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

369 

370 cutoutBytes = packageAlerts.streamCcdDataToBytes(ccdData) 

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

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

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

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

375 

376 ccdData = packageAlerts.createCcdDataCutout( 

377 self.exposure, 

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

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

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

381 self.exposure.getPhotoCalib(), 

382 diaSrcId) 

383 self.assertTrue(ccdData is None) 

384 

385 def testMakeLocalTransformMatrix(self): 

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

387 """ 

388 packageAlerts = PackageAlertsTask() 

389 

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

391 cutout = self.exposure.getCutout(sphPoint, 

392 geom.Extent2I(self.cutoutSize, 

393 self.cutoutSize)) 

394 cd = packageAlerts.makeLocalTransformMatrix( 

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

396 self.assertFloatsAlmostEqual( 

397 cd, 

398 cutout.getWcs().getCdMatrix(), 

399 rtol=5e-10, 

400 atol=5e-10) 

401 

402 def testStreamCcdDataToBytes(self): 

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

404 """ 

405 packageAlerts = PackageAlertsTask() 

406 

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

408 cutout = self.exposure.getCutout(sphPoint, 

409 geom.Extent2I(self.cutoutSize, 

410 self.cutoutSize)) 

411 cutoutCcdData = CCDData( 

412 data=cutout.getImage().array, 

413 wcs=self.cutoutWcs, 

414 unit="adu") 

415 

416 cutoutBytes = packageAlerts.streamCcdDataToBytes(cutoutCcdData) 

417 with io.BytesIO(cutoutBytes) as bytesIO: 

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

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

420 

421 def testFindMissingDiaSourceColumns(self): 

422 """Test that only the DiaSource columns that alerts require are 

423 reported missing. 

424 """ 

425 packageAlerts = PackageAlertsTask() 

426 self.assertEqual(packageAlerts.findMissingDiaSourceColumns(self.diaSources), []) 

427 # templateMjdMean accepts null, so it may be absent. 

428 incomplete = self.diaSources.drop(columns=["mixedInstruments", "templateMjdMean", "templateVersion"]) 

429 self.assertEqual(packageAlerts.findMissingDiaSourceColumns(incomplete), 

430 ["templateVersion", "mixedInstruments"]) 

431 

432 def testMakeAlertDict(self): 

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

434 dictionary "alert". 

435 """ 

436 packageAlerts = PackageAlertsTask() 

437 dia_source_id = 1234 

438 

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

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

441 diaSource["dec"], 

442 geom.degrees) 

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

444 cutout = self.exposure.getCutout(sphPoint, 

445 geom.Extent2I(self.cutoutSize, 

446 self.cutoutSize)) 

447 ccdCutout = packageAlerts.createCcdDataCutout( 

448 cutout, 

449 sphPoint, 

450 pixelPoint, 

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

452 cutout.getPhotoCalib(), 

453 1234, 

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

455 cutoutBytes = packageAlerts.streamCcdDataToBytes( 

456 ccdCutout) 

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

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

459 alert = packageAlerts.makeAlertDict( 

460 dia_source_id, 

461 self.exposure.visitInfo.getObservationReason(), 

462 self.exposure.visitInfo.getObject(), 

463 diaSource, 

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

465 objSources, 

466 objForcedSources, 

467 ccdCutout, 

468 ccdCutout, 

469 ccdCutout) 

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

471 

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

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

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

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

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

477 cutoutBytes) 

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

479 cutoutBytes) 

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

481 cutoutBytes) 

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

483 format="fits") 

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

485 format="fits") 

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

487 template_cutout.header["ROTPA"]) 

488 

489 def testMakeAlertDictSchedulerFields(self): 

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

491 

492 """ 

493 packageAlerts = PackageAlertsTask() 

494 dia_source_id = 1234 

495 

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

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

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

499 obs_reason = f"obs_reason_{srcIdx}", 

500 target = f"target_name_{srcIdx}", 

501 alert = packageAlerts.makeAlertDict( 

502 dia_source_id, 

503 obs_reason, 

504 target, 

505 diaSource, 

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

507 objSources, 

508 objForcedSources, 

509 None, 

510 None, 

511 None) 

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

513 

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

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

516 

517 def testCutoutRotpa(self): 

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

519 """ 

520 packageAlerts = PackageAlertsTask() 

521 

522 # Create a cutout using existing test exposure 

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

524 cutout = self.exposure.getCutout(sphPoint, 

525 geom.Extent2I(self.cutoutSize, 

526 self.cutoutSize)) 

527 

528 # Create CCDData cutout 

529 ccdCutout = packageAlerts.createCcdDataCutout( 

530 cutout, 

531 sphPoint, 

532 self.center, 

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

534 cutout.getPhotoCalib(), 

535 1234, 

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

537 

538 wcs = self.exposure.wcs 

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

540 angle_rad = LocalWcs.computePositionAngle([0], 

541 [cd_matrix[0, 0]], # CD1_1 

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

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

544 [cd_matrix[1, 1]] # CD2_2 

545 ) 

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

547 

548 cutoutBytes = packageAlerts.streamCcdDataToBytes(ccdCutout) 

549 with io.BytesIO(cutoutBytes) as bytesIO: 

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

551 

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

553 self.assertAlmostEqual( 

554 cutoutFromBytes.header['ROTPA'], 

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

556 self.assertAlmostEqual( 

557 cutoutFromBytes.header['ROTPA'], 

558 wcs_pa, places=4) 

559 

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

561 def test_produceAlerts_empty_password(self): 

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

563 """ 

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

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

566 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

567 PackageAlertsTask(config=packConfig) 

568 

569 del self.environ['AP_KAFKA_PRODUCER_PASSWORD'] 

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

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_username(self): 

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

577 """ 

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

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

580 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

581 PackageAlertsTask(config=packConfig) 

582 

583 del self.environ['AP_KAFKA_PRODUCER_USERNAME'] 

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

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_server(self): 

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

591 """ 

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

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

594 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

595 PackageAlertsTask(config=packConfig) 

596 

597 del self.environ['AP_KAFKA_SERVER'] 

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

599 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

600 PackageAlertsTask(config=packConfig) 

601 

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

603 def test_produceAlerts_empty_topic(self): 

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

605 """ 

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

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

608 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

609 PackageAlertsTask(config=packConfig) 

610 

611 del self.environ['AP_KAFKA_TOPIC'] 

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

613 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

614 PackageAlertsTask(config=packConfig) 

615 

616 @patch('confluent_kafka.Producer') 

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

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

619 def test_produceAlerts_success(self, mock_server_check, mock_producer): 

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

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

622 """ 

623 packConfig = PackageAlertsConfig(doProduceAlerts=True) 

624 packageAlerts = PackageAlertsTask(config=packConfig) 

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

626 

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

628 producer_instance = mock_producer.return_value 

629 producer_instance.produce = Mock() 

630 producer_instance.flush = Mock() 

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

632 exposure_time = self.exposure.visitInfo.exposureTime 

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

634 

635 self.assertEqual(mock_server_check.call_count, 1) 

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

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

638 

639 @patch('confluent_kafka.Producer') 

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

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

642 def test_produceAlerts_one_failure(self, mock_server_check, mock_producer): 

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

644 and is writing the failure to disk. 

645 """ 

646 counter = 0 

647 

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

649 nonlocal counter 

650 counter += 1 

651 if counter == 2: 

652 raise KafkaException 

653 else: 

654 return 

655 

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

657 packageAlerts = PackageAlertsTask(config=packConfig) 

658 

659 patcher = patch("builtins.open") 

660 patch_open = patcher.start() 

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

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

663 exposure_time = self.exposure.visitInfo.exposureTime 

664 

665 producer_instance = mock_producer.return_value 

666 producer_instance.produce = Mock(side_effect=mock_produce) 

667 producer_instance.flush = Mock() 

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

669 

670 self.assertEqual(mock_server_check.call_count, 1) 

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

672 self.assertEqual(patch_open.call_count, 1) 

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

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

675 # test above. 

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

677 patcher.stop() 

678 

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

680 def testRun_without_produce(self, mock_server_check): 

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

682 doWriteAlerts set to true. 

683 """ 

684 packConfig = PackageAlertsConfig(doWriteAlerts=True) 

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

686 packConfig.alertWriteLocation = tempdir 

687 packageAlerts = PackageAlertsTask(config=packConfig) 

688 

689 packageAlerts.run(self.diaSources, 

690 self.diaObjects, 

691 self.diaSourceHistory, 

692 self.diaForcedSources, 

693 self.exposure, 

694 self.exposure, 

695 self.exposure) 

696 

697 self.assertEqual(mock_server_check.call_count, 0) 

698 

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

700 writer_schema, data_stream = \ 

701 packageAlerts.alertSchema.retrieve_alerts(f) 

702 data = list(data_stream) 

703 

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

705 for idx, alert in enumerate(data): 

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

707 if isinstance(value, float): 

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

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

710 else: 

711 self.assertAlmostEqual( 

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

713 0.) 

714 elif value is None: 

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

716 else: 

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

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

719 alert["diaSource"]["dec"], 

720 geom.degrees) 

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