Coverage for tests/test_packageAlerts.py: 19%
348 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-11 02:07 -0700
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/>.
22import io
23import os
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
39import lsst.alert.packet as alertPack
40from lsst.ap.association import PackageAlertsConfig, PackageAlertsTask
41from lsst.ap.association.utils import readSchemaFromApdb, convertDataFrameToSdmSchema
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
51import utils_tests
54def _roundTripThroughApdb(objects, sources, forcedSources, dateTime):
55 """Run object and source catalogs through the Apdb to get the correct
56 table schemas.
58 Parameters
59 ----------
60 objects : `pandas.DataFrame`
61 Set of test DiaObjects to round trip.
62 sources : `pandas.DataFrame`
63 Set of test DiaSources to round trip.
64 forcedSources : `pandas.DataFrame`
65 Set of test DiaForcedSources to round trip.
66 dateTime : `astropy.time.Time`
67 Time for the Apdb.
69 Returns
70 -------
71 objects : `pandas.DataFrame`
72 Round tripped objects.
73 sources : `pandas.DataFrame`
74 Round tripped sources.
75 """
76 with tempfile.NamedTemporaryFile() as tmpFile:
77 apdbConfig = ApdbSql.init_database(db_url="sqlite:///" + tmpFile.name)
78 apdb = Apdb.from_config(apdbConfig)
80 wholeSky = Box.full()
81 loadedObjects = apdb.getDiaObjects(wholeSky)
82 if loadedObjects.empty:
83 diaObjects = objects
84 else:
85 diaObjects = pd.concat([loadedObjects, objects])
86 loadedDiaSources = apdb.getDiaSources(wholeSky, [], dateTime)
87 if loadedDiaSources.empty:
88 diaSources = sources
89 else:
90 diaSources = pd.concat([loadedDiaSources, sources])
91 loadedDiaForcedSources = apdb.getDiaForcedSources(wholeSky, [], dateTime)
92 if loadedDiaForcedSources.empty:
93 diaForcedSources = forcedSources
94 else:
95 diaForcedSources = pd.concat([loadedDiaForcedSources, forcedSources])
97 apdb.store(dateTime, diaObjects, diaSources, diaForcedSources)
99 diaObjects = apdb.getDiaObjects(wholeSky)
100 diaSources = apdb.getDiaSources(wholeSky,
101 np.unique(diaObjects["diaObjectId"]),
102 dateTime)
103 diaForcedSources = apdb.getDiaForcedSources(
104 wholeSky, np.unique(diaObjects["diaObjectId"]), dateTime)
106 diaObjects.set_index("diaObjectId", drop=False, inplace=True)
107 diaSources.set_index(["diaObjectId", "band", "diaSourceId"],
108 drop=False,
109 inplace=True)
110 diaForcedSources.set_index(["diaObjectId"], drop=False, inplace=True)
112 # apply SDM type standardization to catch pandas typing issues
113 schema = readSchemaFromApdb(apdb)
114 diaObjects = convertDataFrameToSdmSchema(schema, diaObjects, tableName="DiaObject")
115 diaSources = convertDataFrameToSdmSchema(schema, diaSources, tableName="DiaSource")
116 diaForcedSources = convertDataFrameToSdmSchema(schema, diaForcedSources, tableName="DiaForcedSource")
118 return (diaObjects, diaSources, diaForcedSources)
121VISIT = 2
122DETECTOR = 42
125def mock_alert(dia_source_id):
126 """Generate a minimal mock alert.
127 """
128 return {
129 "diaSourceId": dia_source_id,
130 "diaSource": {
131 "midpointMjdTai": 5,
132 "diaSourceId": 1234,
133 "visit": VISIT,
134 "detector": DETECTOR,
135 "band": 'g',
136 "ra": 12.5,
137 "dec": -16.9,
138 # These types are 32-bit floats in the avro schema, so we have to
139 # make them that type here, so that they round trip appropriately.
140 "x": np.float32(15.7),
141 "y": np.float32(89.8),
142 "apFlux": np.float32(54.85),
143 "apFluxErr": np.float32(70.0),
144 "snr": np.float32(6.7),
145 "psfFlux": np.float32(700.0),
146 "psfFluxErr": np.float32(90.0),
147 # unlike in transformDiaSourceCatalog.py we need a timezone-aware
148 # version because mock_alert does not go through pandas
149 "timeProcessedMjdTai": DateTime.now().get(system=DateTime.MJD, scale=DateTime.TAI)
150 }
151 }
154def mock_ss_alert(dia_source_id, ss_object_id):
155 """Generate a minimal mock alert.
156 """
157 alert = mock_alert(dia_source_id)
158 alert['MPCORB'] = {
159 'mpcDesignation': 'K20A11H', # a string-typed field
160 'ssObjectId': ss_object_id,
161 'q': np.float64(0.99999), # a double-typed field
163 }
164 alert['ssSource'] = {
165 'ssObjectId': ss_object_id,
166 'eclipticLambda': np.float64(3.141592), # a double-typed field
167 'heliocentricDist': np.float32(3.141592), # a float-typed field
168 }
169 return alert
172def _deserialize_alert(alert_bytes):
173 """Deserialize an alert message from Kafka.
175 Parameters
176 ----------
177 alert_bytes : `bytes`
178 Binary-encoding serialized Avro alert, including Confluent Wire
179 Format prefix.
181 Returns
182 -------
183 alert : `dict`
184 An alert payload.
185 """
186 schema = alertPack.Schema.from_uri(str(alertPack.get_uri_to_latest_schema()))
187 content_bytes = io.BytesIO(alert_bytes[5:])
189 return fastavro.schemaless_reader(content_bytes, schema.definition)
192class TestPackageAlerts(lsst.utils.tests.TestCase):
193 def setUp(self):
194 # Create an instance of random generator with fixed seed.
195 rng = np.random.default_rng(1234)
197 patcher = patch.dict(os.environ, {"AP_KAFKA_PRODUCER_PASSWORD": "fake_password",
198 "AP_KAFKA_PRODUCER_USERNAME": "fake_username",
199 "AP_KAFKA_SERVER": "fake_server",
200 "AP_KAFKA_TOPIC": "fake_topic"})
201 self.environ = patcher.start()
202 self.addCleanup(patcher.stop)
203 self.cutoutSize = 35
204 self.center = lsst.geom.Point2D(50.1, 49.8)
205 self.bbox = lsst.geom.Box2I(lsst.geom.Point2I(-20, -30),
206 lsst.geom.Extent2I(140, 160))
207 self.dataset = lsst.meas.base.tests.TestDataset(self.bbox)
208 self.dataset.addSource(100000.0, self.center)
209 exposure, catalog = self.dataset.realize(
210 10.0,
211 self.dataset.makeMinimalSchema(),
212 randomSeed=1234)
213 self.exposure = exposure
214 detector = DetectorWrapper(id=DETECTOR, bbox=exposure.getBBox()).detector
215 self.exposure.setDetector(detector)
217 visit = afwImage.VisitInfo(
218 id=VISIT,
219 exposureTime=200.,
220 date=DateTime("2014-05-13T17:00:00.000000000",
221 DateTime.Timescale.TAI),
222 boresightRotAngle=geom.Angle(0.785398))
223 self.exposure.info.id = 1234
224 self.exposure.info.setVisitInfo(visit)
226 self.exposure.setFilter(
227 afwImage.FilterLabel(band='g', physical="g.MP9401"))
229 diaObjects = utils_tests.makeDiaObjects(2, self.exposure, rng)
230 diaSourceHistory = utils_tests.makeDiaSources(
231 10, diaObjects["diaObjectId"].to_numpy(), self.exposure, rng)
232 diaForcedSources = utils_tests.makeDiaForcedSources(
233 10, diaObjects["diaObjectId"].to_numpy(), self.exposure, rng)
234 self.diaObjects, diaSourceHistory, self.diaForcedSources = _roundTripThroughApdb(
235 diaObjects,
236 diaSourceHistory,
237 diaForcedSources,
238 self.exposure.visitInfo.date.toAstropy())
239 diaSourceHistory["programId"] = 0
241 self.diaSources = diaSourceHistory.loc[[(1, "g", 9), (2, "g", 10)], :]
242 self.diaSources["bboxSize"] = self.cutoutSize
243 self.diaSourceHistory = diaSourceHistory.drop(labels=[(1, "g", 9),
244 (2, "g", 10)])
246 self.cutoutWcs = wcs.WCS(naxis=2)
247 self.cutoutWcs.wcs.crpix = [self.center[0], self.center[1]]
248 self.cutoutWcs.wcs.crval = [
249 self.exposure.getWcs().getSkyOrigin().getRa().asDegrees(),
250 self.exposure.getWcs().getSkyOrigin().getDec().asDegrees()]
251 self.cutoutWcs.wcs.cd = self.exposure.getWcs().getCdMatrix()
252 self.cutoutWcs.wcs.ctype = ["RA---TAN", "DEC--TAN"]
254 def testCreateExtentMinimum(self):
255 """Test the extent creation for the cutout bbox returns a cutout with
256 the minimum cutouut size.
257 """
258 packConfig = PackageAlertsConfig()
259 # Just create a minimum less than the default cutout size.
260 packConfig.minCutoutSize = self.cutoutSize - 5
261 packageAlerts = PackageAlertsTask(config=packConfig)
262 extent = packageAlerts.createDiaSourceExtent(
263 packConfig.minCutoutSize - 5)
264 self.assertTrue(extent == geom.Extent2I(packConfig.minCutoutSize,
265 packConfig.minCutoutSize))
266 # Test that the cutout size is correctly increased.
267 extent = packageAlerts.createDiaSourceExtent(self.cutoutSize)
268 self.assertTrue(extent == geom.Extent2I(self.cutoutSize,
269 self.cutoutSize))
271 def testCreateExtentMaximum(self):
272 """Test the extent creation for the cutout bbox returns a cutout with
273 the maximum cutout size.
274 """
275 packConfig = PackageAlertsConfig()
276 # Just create a maximum more than the default cutout size.
277 packConfig.maxCutoutSize = self.cutoutSize + 5
278 packageAlerts = PackageAlertsTask(config=packConfig)
279 extent = packageAlerts.createDiaSourceExtent(
280 packConfig.maxCutoutSize + 5)
281 self.assertTrue(extent == geom.Extent2I(packConfig.maxCutoutSize,
282 packConfig.maxCutoutSize))
283 # Test that the cutout size is correctly reduced.
284 extent = packageAlerts.createDiaSourceExtent(self.cutoutSize)
285 self.assertTrue(extent == geom.Extent2I(self.cutoutSize,
286 self.cutoutSize))
288 def testCreateCcdDataCutout(self):
289 """Test that the data is being extracted into the CCDData cutout
290 correctly.
291 """
292 packageAlerts = PackageAlertsTask()
294 diaSrcId = 1234
295 ccdData = packageAlerts.createCcdDataCutout(
296 self.exposure,
297 self.exposure.getWcs().getSkyOrigin(),
298 self.exposure.getWcs().getPixelOrigin(),
299 self.exposure.getBBox().getDimensions(),
300 self.exposure.getPhotoCalib(),
301 diaSrcId)
302 calibExposure = self.exposure.getPhotoCalib().calibrateImage(
303 self.exposure.getMaskedImage())
305 self.assertFloatsAlmostEqual(ccdData.wcs.wcs.cd,
306 self.cutoutWcs.wcs.cd)
307 self.assertFloatsAlmostEqual(ccdData.data,
308 calibExposure.getImage().array)
309 self.assertFloatsAlmostEqual(ccdData.psf,
310 self.exposure.psf.computeKernelImage(self.center).array)
312 ccdData = packageAlerts.createCcdDataCutout(
313 self.exposure,
314 geom.SpherePoint(0, 0, geom.degrees),
315 geom.Point2D(-100, -100),
316 self.exposure.getBBox().getDimensions(),
317 self.exposure.getPhotoCalib(),
318 diaSrcId)
319 self.assertTrue(ccdData is None)
321 def testMakeLocalTransformMatrix(self):
322 """Test that the local WCS approximation is correct.
323 """
324 packageAlerts = PackageAlertsTask()
326 sphPoint = self.exposure.getWcs().pixelToSky(self.center)
327 cutout = self.exposure.getCutout(sphPoint,
328 geom.Extent2I(self.cutoutSize,
329 self.cutoutSize))
330 cd = packageAlerts.makeLocalTransformMatrix(
331 cutout.getWcs(), self.center, sphPoint)
332 self.assertFloatsAlmostEqual(
333 cd,
334 cutout.getWcs().getCdMatrix(),
335 rtol=5e-10,
336 atol=5e-10)
338 def testStreamCcdDataToBytes(self):
339 """Test round tripping an CCDData cutout to bytes and back.
340 """
341 packageAlerts = PackageAlertsTask()
343 sphPoint = self.exposure.getWcs().pixelToSky(self.center)
344 cutout = self.exposure.getCutout(sphPoint,
345 geom.Extent2I(self.cutoutSize,
346 self.cutoutSize))
347 cutoutCcdData = CCDData(
348 data=cutout.getImage().array,
349 wcs=self.cutoutWcs,
350 unit="adu")
352 cutoutBytes = packageAlerts.streamCcdDataToBytes(cutoutCcdData)
353 with io.BytesIO(cutoutBytes) as bytesIO:
354 cutoutFromBytes = CCDData.read(bytesIO, format="fits")
355 self.assertFloatsAlmostEqual(cutoutCcdData.data, cutoutFromBytes.data)
357 def testMakeAlertDict(self):
358 """Test stripping data from the various data products and into a
359 dictionary "alert".
360 """
361 packageAlerts = PackageAlertsTask()
362 dia_source_id = 1234
364 for srcIdx, diaSource in self.diaSources.iterrows():
365 sphPoint = geom.SpherePoint(diaSource["ra"],
366 diaSource["dec"],
367 geom.degrees)
368 pixelPoint = geom.Point2D(diaSource["x"], diaSource["y"])
369 cutout = self.exposure.getCutout(sphPoint,
370 geom.Extent2I(self.cutoutSize,
371 self.cutoutSize))
372 ccdCutout = packageAlerts.createCcdDataCutout(
373 cutout,
374 sphPoint,
375 pixelPoint,
376 geom.Extent2I(self.cutoutSize, self.cutoutSize),
377 cutout.getPhotoCalib(),
378 1234,
379 rotPa=cutout.visitInfo.boresightRotAngle.asDegrees())
380 cutoutBytes = packageAlerts.streamCcdDataToBytes(
381 ccdCutout)
382 objSources = self.diaSourceHistory.loc[srcIdx[0]]
383 objForcedSources = self.diaForcedSources.loc[srcIdx[0]]
384 alert = packageAlerts.makeAlertDict(
385 dia_source_id,
386 diaSource,
387 self.diaObjects.loc[srcIdx[0]],
388 objSources,
389 objForcedSources,
390 ccdCutout,
391 ccdCutout,
392 ccdCutout)
393 self.assertEqual(len(alert), 11)
395 self.assertEqual(alert["diaSourceId"], dia_source_id)
396 self.assertEqual(alert["diaSource"], diaSource.to_dict())
397 self.assertEqual(alert["cutoutDifference"],
398 cutoutBytes)
399 self.assertEqual(alert["cutoutScience"],
400 cutoutBytes)
401 self.assertEqual(alert["cutoutTemplate"],
402 cutoutBytes)
403 science_cutout = CCDData.read(io.BytesIO(alert["cutoutScience"]),
404 format="fits")
405 template_cutout = CCDData.read(io.BytesIO(alert["cutoutTemplate"]),
406 format="fits")
407 self.assertAlmostEqual(science_cutout.header["ROTPA"],
408 template_cutout.header["ROTPA"])
410 def testCutoutRotpa(self):
411 """Test that the ROTPA header keyword matches the boresightRotAngle from visitInfo.
412 """
413 packageAlerts = PackageAlertsTask()
415 # Create a cutout using existing test exposure
416 sphPoint = self.exposure.getWcs().pixelToSky(self.center)
417 cutout = self.exposure.getCutout(sphPoint,
418 geom.Extent2I(self.cutoutSize,
419 self.cutoutSize))
421 # Create CCDData cutout
422 ccdCutout = packageAlerts.createCcdDataCutout(
423 cutout,
424 sphPoint,
425 self.center,
426 geom.Extent2I(self.cutoutSize, self.cutoutSize),
427 cutout.getPhotoCalib(),
428 1234,
429 rotPa=cutout.visitInfo.boresightRotAngle.asDegrees())
431 wcs = self.exposure.wcs
432 cd_matrix = wcs.getCdMatrix() # This gets the CD matrix elements
433 angle_rad = LocalWcs.computePositionAngle([0],
434 [cd_matrix[0, 0]], # CD1_1
435 [cd_matrix[0, 1]], # CD1_2
436 [cd_matrix[1, 0]], # CD2_1
437 [cd_matrix[1, 1]] # CD2_2
438 )
439 wcs_pa = geom.Angle(angle_rad, geom.radians).asDegrees()
441 cutoutBytes = packageAlerts.streamCcdDataToBytes(ccdCutout)
442 with io.BytesIO(cutoutBytes) as bytesIO:
443 cutoutFromBytes = CCDData.read(bytesIO, format="fits")
445 self.assertIn('ROTPA', cutoutFromBytes.header)
446 self.assertAlmostEqual(
447 cutoutFromBytes.header['ROTPA'],
448 self.exposure.visitInfo.boresightRotAngle.asDegrees())
449 self.assertAlmostEqual(
450 cutoutFromBytes.header['ROTPA'],
451 wcs_pa, places=4)
453 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled')
454 def test_produceAlerts_empty_password(self):
455 """ Test that produceAlerts raises if the password is empty or missing.
456 """
457 self.environ['AP_KAFKA_PRODUCER_PASSWORD'] = ""
458 with self.assertRaisesRegex(ValueError, "Kafka password"):
459 packConfig = PackageAlertsConfig(doProduceAlerts=True)
460 PackageAlertsTask(config=packConfig)
462 del self.environ['AP_KAFKA_PRODUCER_PASSWORD']
463 with self.assertRaisesRegex(ValueError, "Kafka password"):
464 packConfig = PackageAlertsConfig(doProduceAlerts=True)
465 PackageAlertsTask(config=packConfig)
467 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled')
468 def test_produceAlerts_empty_username(self):
469 """ Test that produceAlerts raises if the username is empty or missing.
470 """
471 self.environ['AP_KAFKA_PRODUCER_USERNAME'] = ""
472 with self.assertRaisesRegex(ValueError, "Kafka username"):
473 packConfig = PackageAlertsConfig(doProduceAlerts=True)
474 PackageAlertsTask(config=packConfig)
476 del self.environ['AP_KAFKA_PRODUCER_USERNAME']
477 with self.assertRaisesRegex(ValueError, "Kafka username"):
478 packConfig = PackageAlertsConfig(doProduceAlerts=True)
479 PackageAlertsTask(config=packConfig)
481 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled')
482 def test_produceAlerts_empty_server(self):
483 """ Test that produceAlerts raises if the server is empty or missing.
484 """
485 self.environ['AP_KAFKA_SERVER'] = ""
486 with self.assertRaisesRegex(ValueError, "Kafka server"):
487 packConfig = PackageAlertsConfig(doProduceAlerts=True)
488 PackageAlertsTask(config=packConfig)
490 del self.environ['AP_KAFKA_SERVER']
491 with self.assertRaisesRegex(ValueError, "Kafka server"):
492 packConfig = PackageAlertsConfig(doProduceAlerts=True)
493 PackageAlertsTask(config=packConfig)
495 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled')
496 def test_produceAlerts_empty_topic(self):
497 """ Test that produceAlerts raises if the topic is empty or missing.
498 """
499 self.environ['AP_KAFKA_TOPIC'] = ""
500 with self.assertRaisesRegex(ValueError, "Kafka topic"):
501 packConfig = PackageAlertsConfig(doProduceAlerts=True)
502 PackageAlertsTask(config=packConfig)
504 del self.environ['AP_KAFKA_TOPIC']
505 with self.assertRaisesRegex(ValueError, "Kafka topic"):
506 packConfig = PackageAlertsConfig(doProduceAlerts=True)
507 PackageAlertsTask(config=packConfig)
509 @patch('confluent_kafka.Producer')
510 @patch.object(PackageAlertsTask, '_server_check')
511 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled')
512 def test_produceAlerts_success(self, mock_server_check, mock_producer):
513 """ Test that produceAlerts calls the producer on all provided alerts
514 when the alerts are all under the batch size limit.
515 """
516 packConfig = PackageAlertsConfig(doProduceAlerts=True)
517 packageAlerts = PackageAlertsTask(config=packConfig)
518 alerts = [mock_alert(1), mock_alert(2), mock_ss_alert(3, 3)]
520 # Create a variable and assign it an instance of the patched kafka producer
521 producer_instance = mock_producer.return_value
522 producer_instance.produce = Mock()
523 producer_instance.flush = Mock()
524 unix_midpoint = self.exposure.visitInfo.date.toAstropy().tai.unix
525 exposure_time = self.exposure.visitInfo.exposureTime
526 packageAlerts.produceAlerts(alerts, VISIT, DETECTOR, unix_midpoint, exposure_time)
528 self.assertEqual(mock_server_check.call_count, 1)
529 self.assertEqual(producer_instance.produce.call_count, len(alerts))
530 self.assertEqual(producer_instance.flush.call_count, len(alerts)+1)
532 @patch('confluent_kafka.Producer')
533 @patch.object(PackageAlertsTask, '_server_check')
534 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled')
535 def test_produceAlerts_one_failure(self, mock_server_check, mock_producer):
536 """ Test that produceAlerts correctly fails on one alert
537 and is writing the failure to disk.
538 """
539 counter = 0
541 def mock_produce(*args, **kwargs):
542 nonlocal counter
543 counter += 1
544 if counter == 2:
545 raise KafkaException
546 else:
547 return
549 packConfig = PackageAlertsConfig(doProduceAlerts=True, doWriteFailedAlerts=True)
550 packageAlerts = PackageAlertsTask(config=packConfig)
552 patcher = patch("builtins.open")
553 patch_open = patcher.start()
554 alerts = [mock_alert(1), mock_alert(2), mock_alert(3), mock_ss_alert(4, 4)]
555 unix_midpoint = self.exposure.visitInfo.date.toAstropy().tai.unix
556 exposure_time = self.exposure.visitInfo.exposureTime
558 producer_instance = mock_producer.return_value
559 producer_instance.produce = Mock(side_effect=mock_produce)
560 producer_instance.flush = Mock()
561 packageAlerts.produceAlerts(alerts, VISIT, DETECTOR, unix_midpoint, exposure_time)
563 self.assertEqual(mock_server_check.call_count, 1)
564 self.assertEqual(producer_instance.produce.call_count, len(alerts))
565 self.assertEqual(patch_open.call_count, 1)
566 self.assertIn(f"{VISIT}_{DETECTOR}_2.avro", patch_open.call_args.args[0])
567 # Because one produce raises, we call flush one fewer times than in the success
568 # test above.
569 self.assertEqual(producer_instance.flush.call_count, len(alerts))
570 patcher.stop()
572 @patch.object(PackageAlertsTask, '_server_check')
573 def testRun_without_produce(self, mock_server_check):
574 """Test the run method of package alerts with produce set to False and
575 doWriteAlerts set to true.
576 """
577 packConfig = PackageAlertsConfig(doWriteAlerts=True)
578 with tempfile.TemporaryDirectory(prefix='alerts') as tempdir:
579 packConfig.alertWriteLocation = tempdir
580 packageAlerts = PackageAlertsTask(config=packConfig)
582 packageAlerts.run(self.diaSources,
583 self.diaObjects,
584 self.diaSourceHistory,
585 self.diaForcedSources,
586 self.exposure,
587 self.exposure,
588 self.exposure)
590 self.assertEqual(mock_server_check.call_count, 0)
592 with open(os.path.join(tempdir, f"{VISIT}_{DETECTOR}.avro"), 'rb') as f:
593 writer_schema, data_stream = \
594 packageAlerts.alertSchema.retrieve_alerts(f)
595 data = list(data_stream)
597 self.assertEqual(len(data), len(self.diaSources))
598 for idx, alert in enumerate(data):
599 for key, value in alert["diaSource"].items():
600 if isinstance(value, float):
601 if np.isnan(self.diaSources.iloc[idx][key]):
602 self.assertTrue(np.isnan(value))
603 else:
604 self.assertAlmostEqual(
605 1 - value / self.diaSources.iloc[idx][key],
606 0.)
607 elif value is None:
608 self.assertTrue(pd.isna(self.diaSources.iloc[idx][key]))
609 else:
610 self.assertEqual(value, self.diaSources.iloc[idx][key])
611 sphPoint = geom.SpherePoint(alert["diaSource"]["ra"],
612 alert["diaSource"]["dec"],
613 geom.degrees)
614 pixelPoint = geom.Point2D(alert["diaSource"]["x"], alert["diaSource"]["y"])
615 cutout = self.exposure.getCutout(sphPoint,
616 geom.Extent2I(self.cutoutSize,
617 self.cutoutSize))
618 ccdCutout = packageAlerts.createCcdDataCutout(
619 cutout,
620 sphPoint,
621 pixelPoint,
622 geom.Extent2I(self.cutoutSize, self.cutoutSize),
623 cutout.getPhotoCalib(),
624 1234,
625 rotPa=cutout.visitInfo.boresightRotAngle.asDegrees())
627 self.assertEqual(alert["cutoutDifference"],
628 packageAlerts.streamCcdDataToBytes(ccdCutout))
630 self.assertEqual(alert["cutoutDifference"],
631 packageAlerts.streamCcdDataToBytes(ccdCutout))
633 @patch.object(PackageAlertsTask, '_server_check')
634 def testRun_without_produce_use_averagePsf(self, mock_server_check):
635 """Test the run method of package alerts with produce set to False and
636 doWriteAlerts set to true.
637 """
638 packConfig = PackageAlertsConfig(doWriteAlerts=True)
639 with tempfile.TemporaryDirectory(prefix='alerts') as tempdir:
640 packConfig.alertWriteLocation = tempdir
641 packConfig.useAveragePsf = True
642 packageAlerts = PackageAlertsTask(config=packConfig)
644 packageAlerts.run(self.diaSources,
645 self.diaObjects,
646 self.diaSourceHistory,
647 self.diaForcedSources,
648 self.exposure,
649 self.exposure,
650 self.exposure)
652 self.assertEqual(mock_server_check.call_count, 0)
654 with open(os.path.join(tempdir, f"{VISIT}_{DETECTOR}.avro"), 'rb') as f:
655 writer_schema, data_stream = \
656 packageAlerts.alertSchema.retrieve_alerts(f)
657 data = list(data_stream)
659 self.assertEqual(len(data), len(self.diaSources))
660 for idx, alert in enumerate(data):
661 for key, value in alert["diaSource"].items():
662 if isinstance(value, float):
663 if np.isnan(self.diaSources.iloc[idx][key]):
664 self.assertTrue(np.isnan(value))
665 else:
666 self.assertAlmostEqual(
667 1 - value / self.diaSources.iloc[idx][key],
668 0.)
669 elif value is None:
670 self.assertTrue(pd.isna(self.diaSources.iloc[idx][key]))
671 else:
672 self.assertEqual(value, self.diaSources.iloc[idx][key])
673 sphPoint = geom.SpherePoint(alert["diaSource"]["ra"],
674 alert["diaSource"]["dec"],
675 geom.degrees)
676 pixelPoint = geom.Point2D(alert["diaSource"]["x"], alert["diaSource"]["y"])
677 cutout = self.exposure.getCutout(sphPoint,
678 geom.Extent2I(self.cutoutSize,
679 self.cutoutSize))
680 ccdCutout = packageAlerts.createCcdDataCutout(
681 cutout,
682 sphPoint,
683 pixelPoint,
684 geom.Extent2I(self.cutoutSize, self.cutoutSize),
685 cutout.getPhotoCalib(),
686 1234,
687 rotPa=cutout.visitInfo.boresightRotAngle.asDegrees())
688 self.assertEqual(alert["cutoutDifference"],
689 packageAlerts.streamCcdDataToBytes(ccdCutout))
691 @patch.object(PackageAlertsTask, 'produceAlerts')
692 @patch('confluent_kafka.Producer')
693 @patch.object(PackageAlertsTask, '_server_check')
694 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled')
695 def testRun_with_produce(self, mock_produceAlerts, mock_server_check, mock_producer):
696 """Test that packageAlerts calls produceAlerts when doProduceAlerts
697 is set to True.
698 """
699 packConfig = PackageAlertsConfig(doProduceAlerts=True)
700 packageAlerts = PackageAlertsTask(config=packConfig)
702 packageAlerts.run(self.diaSources,
703 self.diaObjects,
704 self.diaSourceHistory,
705 self.diaForcedSources,
706 self.exposure,
707 self.exposure,
708 self.exposure)
709 self.assertEqual(mock_server_check.call_count, 1)
710 self.assertEqual(mock_produceAlerts.call_count, 1)
712 def test_serialize_alert_round_trip(self):
713 """Test that values in the alert packet exactly round trip.
714 """
715 packClass = PackageAlertsConfig()
716 packageAlerts = PackageAlertsTask(config=packClass)
718 alert = mock_ss_alert(1, 1)
719 serialized = PackageAlertsTask._serializeAlert(packageAlerts, alert)
720 deserialized = _deserialize_alert(serialized)
721 for table in ['diaSource', 'ssSource', 'MPCORB']:
722 for field in alert[table]:
723 self.assertEqual(alert[table][field], deserialized[table][field])
725 self.assertEqual(1, deserialized["diaSourceId"])
727 @unittest.skipIf(confluent_kafka is None, 'Kafka is not enabled')
728 def test_server_check(self):
730 with self.assertRaisesRegex(KafkaException, "_TRANSPORT"):
731 packConfig = PackageAlertsConfig(doProduceAlerts=True)
732 PackageAlertsTask(config=packConfig)
735class MemoryTester(lsst.utils.tests.MemoryTestCase):
736 pass
739def setup_module(module):
740 lsst.utils.tests.init()
743if __name__ == "__main__": 743 ↛ 744line 743 didn't jump to line 744 because the condition on line 743 was never true
744 lsst.utils.tests.init()
745 unittest.main()