Coverage for tests/test_ppdb.py: 99%
306 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-22 02:18 -0700
1# This file is part of analysis_ap.
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 os
23import unittest
24import unittest.mock as mock
26import astropy.table
28import lsst.afw.geom
29import lsst.afw.image
30import lsst.geom
31import lsst.utils.tests
32from lsst.analysis.ap.ppdb import (
33 DEFAULT_PPDB_TAP_URL,
34 DiaObjectLightCurve,
35 PpdbTap,
36 region_from_exposure,
37)
40class _FakeTapResults:
41 """Stand-in for a pyvo TAPResults; just wraps a fixed astropy Table."""
43 def __init__(self, table):
44 self._table = table
46 def to_table(self):
47 return self._table
50class _FakeTapService:
51 """Record ADQL queries and return preset tables, without any network.
53 ``tables`` may be a single Table (returned for every call) or a list of
54 Tables returned for successive calls (the last is repeated thereafter).
55 """
57 def __init__(self, tables):
58 if isinstance(tables, astropy.table.Table):
59 tables = [tables]
60 self._tables = list(tables)
61 self.queries = []
62 self.maxrecs = []
64 def run_async(self, query, maxrec=None):
65 self.queries.append(query)
66 self.maxrecs.append(maxrec)
67 index = min(len(self.queries) - 1, len(self._tables) - 1)
68 return _FakeTapResults(self._tables[index])
71def _objects_table(dia_object_ids=(1, 2, 3)):
72 n = len(dia_object_ids)
73 return astropy.table.Table({
74 "diaObjectId": list(dia_object_ids),
75 "ra": [150.0 + 0.1 * i for i in range(n)],
76 "dec": [2.5 + 0.1 * i for i in range(n)],
77 })
80def _sources_table(n=2, id_column="diaSourceId"):
81 return astropy.table.Table({
82 id_column: [100 + i for i in range(n)],
83 "diaObjectId": [42] * n,
84 "midpointMjdTai": [60000.0 + i for i in range(n)],
85 "band": ["r"] * n,
86 })
89def _make_exposure(ra=150.0, dec=2.5, scale=0.2, size=100):
90 """A square ExposureF with a simple TAN WCS centered at (ra, dec)."""
91 bbox = lsst.geom.Box2I(lsst.geom.Point2I(0, 0),
92 lsst.geom.Extent2I(size, size))
93 exposure = lsst.afw.image.ExposureF(bbox)
94 crpix = lsst.geom.Point2D(size / 2.0, size / 2.0)
95 crval = lsst.geom.SpherePoint(ra, dec, lsst.geom.degrees)
96 cd_matrix = lsst.afw.geom.makeCdMatrix(scale=scale * lsst.geom.arcseconds)
97 exposure.setWcs(lsst.afw.geom.makeSkyWcs(crpix, crval, cd_matrix))
98 return exposure
101class TestRegionFromExposure(lsst.utils.tests.TestCase):
102 def test_center_and_radius(self):
103 exposure = _make_exposure(ra=150.0, dec=2.5, scale=0.2, size=100)
104 ra, dec, radius0 = region_from_exposure(exposure, padding=0.0)
105 self.assertAlmostEqual(ra, 150.0, places=3)
106 self.assertAlmostEqual(dec, 2.5, places=3)
107 # ~half-diagonal of a 100px box at 0.2 arcsec/px is ~14 arcsec.
108 self.assertAlmostEqual(radius0 * 3600.0, 14.0, delta=0.5)
110 def test_padding_is_arcseconds(self):
111 exposure = _make_exposure()
112 _, _, radius0 = region_from_exposure(exposure, padding=0.0)
113 _, _, radius5 = region_from_exposure(exposure, padding=5.0)
114 # padding must add exactly 5 arcsec to the radius.
115 self.assertAlmostEqual((radius5 - radius0) * 3600.0, 5.0, places=6)
117 def test_no_wcs_raises(self):
118 bbox = lsst.geom.Box2I(lsst.geom.Point2I(0, 0),
119 lsst.geom.Extent2I(10, 10))
120 exposure = lsst.afw.image.ExposureF(bbox)
121 with self.assertRaisesRegex(ValueError, "no WCS"):
122 region_from_exposure(exposure)
124 def test_negative_padding_raises(self):
125 with self.assertRaisesRegex(ValueError, "non-negative"):
126 region_from_exposure(_make_exposure(), padding=-1.0)
129class TestPpdbTapAuth(lsst.utils.tests.TestCase):
130 def test_missing_token_raises(self):
131 with mock.patch.dict(os.environ, {}, clear=True):
132 with self.assertRaisesRegex(RuntimeError, "No RSP token"):
133 PpdbTap()
135 def test_token_sets_bearer_header(self):
136 with mock.patch("lsst.analysis.ap.ppdb.pyvo.dal.TAPService") as tap:
137 PpdbTap(token="secret-token")
138 args, kwargs = tap.call_args
139 self.assertEqual(args[0], DEFAULT_PPDB_TAP_URL)
140 session = kwargs["session"]
141 self.assertEqual(session.headers["Authorization"],
142 "Bearer secret-token")
144 def test_env_token_used(self):
145 with mock.patch.dict(os.environ, {"RSP_TOKEN": "from-env"}, clear=True):
146 with mock.patch("lsst.analysis.ap.ppdb.pyvo.dal.TAPService") as tap:
147 PpdbTap()
148 session = tap.call_args.kwargs["session"]
149 self.assertEqual(session.headers["Authorization"], "Bearer from-env")
152class TestLoadObjects(lsst.utils.tests.TestCase):
153 def test_default_query(self):
154 fake = _FakeTapService(_objects_table())
155 PpdbTap(service=fake).load_objects()
156 query = fake.queries[-1]
157 self.assertIn("FROM ppdb.DiaObject", query)
158 self.assertIn("validityEndMjdTai IS NULL", query)
159 self.assertIn("TOP 100000", query)
160 self.assertIn("ORDER BY diaObjectId", query)
162 def test_latest_false_drops_validity_filter(self):
163 fake = _FakeTapService(_objects_table())
164 PpdbTap(service=fake).load_objects(latest=False)
165 self.assertNotIn("validityEndMjdTai", fake.queries[-1])
167 def test_cone_search(self):
168 fake = _FakeTapService(_objects_table())
169 PpdbTap(service=fake).load_objects(ra=150.0, dec=2.5, radius=1.0)
170 query = fake.queries[-1]
171 self.assertIn("CONTAINS(POINT('ICRS', ra, dec)", query)
172 self.assertIn("CIRCLE('ICRS', 150.0, 2.5, 1.0)", query)
174 def test_partial_cone_raises(self):
175 fake = _FakeTapService(_objects_table())
176 with self.assertRaisesRegex(ValueError, "ra, dec, and radius"):
177 PpdbTap(service=fake).load_objects(ra=150.0, dec=2.5)
179 def test_exposure_region(self):
180 fake = _FakeTapService(_objects_table())
181 PpdbTap(service=fake).load_objects(exposure=_make_exposure())
182 self.assertIn("CONTAINS(POINT('ICRS', ra, dec)", fake.queries[-1])
184 def test_exposure_and_cone_conflict(self):
185 fake = _FakeTapService(_objects_table())
186 with self.assertRaisesRegex(ValueError, "either exposure="):
187 PpdbTap(service=fake).load_objects(exposure=_make_exposure(),
188 ra=150.0, dec=2.5, radius=1.0)
190 def test_no_limit_omits_top(self):
191 fake = _FakeTapService(_objects_table())
192 PpdbTap(service=fake).load_objects(limit=None)
193 self.assertNotIn("TOP", fake.queries[-1])
195 def test_duplicate_current_versions_warns(self):
196 fake = _FakeTapService(_objects_table((1, 1, 2)))
197 ppdb = PpdbTap(service=fake)
198 with self.assertLogs("lsst.analysis.ap.ppdb", level="WARNING") as cm:
199 ppdb.load_objects()
200 self.assertTrue(any("duplicate current versions" in m
201 for m in cm.output))
203 def test_truncation_warns(self):
204 fake = _FakeTapService(_objects_table((1, 2)))
205 ppdb = PpdbTap(service=fake)
206 with self.assertLogs("lsst.analysis.ap.ppdb", level="WARNING") as cm:
207 ppdb.load_objects(limit=2)
208 self.assertTrue(any("row limit" in m for m in cm.output))
210 def test_load_object_missing_raises(self):
211 fake = _FakeTapService(astropy.table.Table())
212 with self.assertRaisesRegex(RuntimeError, "diaObjectId=999 not found"):
213 PpdbTap(service=fake).load_object(999)
215 def test_load_object_returns_row(self):
216 fake = _FakeTapService(_objects_table((7,)))
217 row = PpdbTap(service=fake).load_object(7)
218 self.assertEqual(row["diaObjectId"], 7)
219 self.assertIn("diaObjectId = 7", fake.queries[-1])
220 self.assertIn("validityEndMjdTai IS NULL", fake.queries[-1])
223class TestLoadSources(lsst.utils.tests.TestCase):
224 def test_sources_for_object(self):
225 fake = _FakeTapService(_sources_table())
226 PpdbTap(service=fake).load_sources_for_object(42)
227 query = fake.queries[-1]
228 self.assertIn("FROM ppdb.DiaSource", query)
229 self.assertIn("diaObjectId = 42", query)
230 self.assertIn("ORDER BY midpointMjdTai", query)
232 def test_forced_sources_for_object(self):
233 fake = _FakeTapService(_sources_table(id_column="diaForcedSourceId"))
234 PpdbTap(service=fake).load_forced_sources_for_object(42)
235 self.assertIn("FROM ppdb.DiaForcedSource", fake.queries[-1])
237 def test_band_filter(self):
238 fake = _FakeTapService(_sources_table())
239 PpdbTap(service=fake).load_sources(diaObjectId=42, bands=["r", "g"])
240 query = fake.queries[-1]
241 self.assertIn("diaObjectId = 42", query)
242 self.assertIn("band IN ('r', 'g')", query)
244 def test_invalid_band_raises(self):
245 fake = _FakeTapService(_sources_table())
246 with self.assertRaisesRegex(ValueError, "Unknown band"):
247 PpdbTap(service=fake).load_sources(diaObjectId=42, bands=["x"])
249 def test_time_window(self):
250 fake = _FakeTapService(_sources_table())
251 PpdbTap(service=fake).load_sources(diaObjectId=42, mjd_begin=100.0,
252 mjd_end=200.0)
253 query = fake.queries[-1]
254 self.assertIn("midpointMjdTai >= 100.0", query)
255 self.assertIn("midpointMjdTai < 200.0", query)
257 def test_object_id_required(self):
258 fake = _FakeTapService(_sources_table())
259 # omitting it entirely is a TypeError (required keyword)
260 with self.assertRaises(TypeError):
261 PpdbTap(service=fake).load_sources()
262 # passing None explicitly gives a guiding ValueError
263 with self.assertRaisesRegex(ValueError, "diaObjectId is required"):
264 PpdbTap(service=fake).load_sources(diaObjectId=None)
266 def test_load_source_missing_raises(self):
267 fake = _FakeTapService(astropy.table.Table())
268 with self.assertRaisesRegex(RuntimeError, "diaSourceId=5 not found"):
269 PpdbTap(service=fake).load_source(5)
271 def test_load_source_returns_row(self):
272 fake = _FakeTapService(_sources_table(n=1))
273 row = PpdbTap(service=fake).load_source(100)
274 self.assertEqual(row["diaSourceId"], 100)
276 def test_id_list_is_chunked(self):
277 fake = _FakeTapService(_sources_table(n=2))
278 result = PpdbTap(service=fake).load_sources(
279 diaObjectId=[1, 2, 3, 4, 5], id_chunk_size=2)
280 # 5 ids in chunks of 2 -> 3 IN queries, results concatenated.
281 self.assertEqual(len(fake.queries), 3)
282 self.assertIn("diaObjectId IN (1, 2)", fake.queries[0])
283 self.assertIn("diaObjectId IN (3, 4)", fake.queries[1])
284 self.assertIn("diaObjectId IN (5)", fake.queries[2])
285 self.assertEqual(len(result), 6)
287 def test_empty_id_list_returns_empty(self):
288 fake = _FakeTapService(_sources_table(n=0))
289 result = PpdbTap(service=fake).load_sources(diaObjectId=[])
290 self.assertEqual(len(result), 0)
291 self.assertIn("1 = 0", fake.queries[-1])
294class TestRegionAndCone(lsst.utils.tests.TestCase):
295 def test_sources_for_region_two_step(self):
296 fake = _FakeTapService([_objects_table((10, 20, 30)),
297 _sources_table(n=2)])
298 PpdbTap(service=fake).load_sources_for_region(ra=150.0, dec=2.5,
299 radius=0.1)
300 # first: cone on DiaObject (the only table allowing spatial search).
301 obj_query = fake.queries[0]
302 self.assertIn("FROM ppdb.DiaObject", obj_query)
303 self.assertIn("validityEndMjdTai IS NULL", obj_query)
304 self.assertIn("CONTAINS(POINT('ICRS', ra, dec)", obj_query)
305 # second: sources by id, with no spatial predicate.
306 src_query = fake.queries[1]
307 self.assertIn("FROM ppdb.DiaSource", src_query)
308 self.assertIn("diaObjectId IN (10, 20, 30)", src_query)
309 self.assertNotIn("CONTAINS", src_query)
311 def test_forced_sources_for_region_with_exposure(self):
312 fake = _FakeTapService([
313 _objects_table((7,)),
314 _sources_table(n=1, id_column="diaForcedSourceId"),
315 ])
316 PpdbTap(service=fake).load_forced_sources_for_region(
317 exposure=_make_exposure())
318 self.assertIn("FROM ppdb.DiaObject", fake.queries[0])
319 self.assertIn("CONTAINS", fake.queries[0])
320 self.assertIn("FROM ppdb.DiaForcedSource", fake.queries[1])
321 self.assertIn("diaObjectId IN (7)", fake.queries[1])
323 def test_region_with_no_objects(self):
324 fake = _FakeTapService([astropy.table.Table(), _sources_table(n=0)])
325 result = PpdbTap(service=fake).load_sources_for_region(
326 ra=1.0, dec=1.0, radius=0.01)
327 self.assertEqual(len(result), 0)
328 self.assertIn("1 = 0", fake.queries[1])
330 def test_sources_by_cone_warns_and_queries_source_table(self):
331 fake = _FakeTapService(_sources_table(n=2))
332 ppdb = PpdbTap(service=fake)
333 with self.assertLogs("lsst.analysis.ap.ppdb", level="WARNING") as cm:
334 ppdb.load_sources_by_cone(ra=150.0, dec=2.5, radius=0.05)
335 self.assertTrue(any("prototype-only" in m for m in cm.output))
336 query = fake.queries[-1]
337 self.assertIn("FROM ppdb.DiaSource", query)
338 self.assertIn("CONTAINS(POINT('ICRS', ra, dec)", query)
340 def test_forced_sources_by_cone(self):
341 fake = _FakeTapService(
342 _sources_table(n=1, id_column="diaForcedSourceId"))
343 ppdb = PpdbTap(service=fake)
344 with self.assertLogs("lsst.analysis.ap.ppdb", level="WARNING"):
345 ppdb.load_forced_sources_by_cone(exposure=_make_exposure())
346 self.assertIn("FROM ppdb.DiaForcedSource", fake.queries[-1])
347 self.assertIn("CONTAINS", fake.queries[-1])
349 def test_by_cone_requires_a_cone(self):
350 fake = _FakeTapService(_sources_table())
351 with self.assertRaisesRegex(ValueError, "cone search requires"):
352 PpdbTap(service=fake).load_sources_by_cone()
354 def test_sources_by_time_window_warns_and_queries_source_table(self):
355 fake = _FakeTapService(_sources_table(n=2))
356 ppdb = PpdbTap(service=fake)
357 with self.assertLogs("lsst.analysis.ap.ppdb", level="WARNING") as cm:
358 ppdb.load_sources_by_time_window(mjd_begin=60000.0, mjd_end=60001.0)
359 self.assertTrue(any("prototype-only" in m for m in cm.output))
360 query = fake.queries[-1]
361 self.assertIn("FROM ppdb.DiaSource", query)
362 self.assertIn("midpointMjdTai >= 60000.0", query)
363 self.assertIn("midpointMjdTai < 60001.0", query)
364 self.assertNotIn("diaObjectId", query)
365 self.assertNotIn("CONTAINS", query)
367 def test_forced_sources_by_time_window_open_ended(self):
368 fake = _FakeTapService(
369 _sources_table(n=1, id_column="diaForcedSourceId"))
370 ppdb = PpdbTap(service=fake)
371 with self.assertLogs("lsst.analysis.ap.ppdb", level="WARNING"):
372 ppdb.load_forced_sources_by_time_window(mjd_begin=60000.0)
373 query = fake.queries[-1]
374 self.assertIn("FROM ppdb.DiaForcedSource", query)
375 self.assertIn("midpointMjdTai >= 60000.0", query)
377 def test_by_time_window_requires_a_bound(self):
378 fake = _FakeTapService(_sources_table())
379 with self.assertRaisesRegex(ValueError, "time-window search requires"):
380 PpdbTap(service=fake).load_sources_by_time_window()
383class TestColumnOrdering(lsst.utils.tests.TestCase):
384 def test_sources_columns_reordered_to_sdm(self):
385 # alphabetical input columns -> SDM-schema order out.
386 scrambled = astropy.table.Table({
387 "band": ["r"], "dec": [1.0], "diaObjectId": [42],
388 "diaSourceId": [100], "midpointMjdTai": [60000.0], "ra": [150.0]})
389 fake = _FakeTapService(scrambled)
390 result = PpdbTap(service=fake).load_sources(diaObjectId=42)
391 self.assertEqual(result.colnames,
392 ["diaSourceId", "diaObjectId", "midpointMjdTai",
393 "ra", "dec", "band"])
395 def test_unknown_columns_raise(self):
396 # a column not in the SDM schema signals schema drift -> error.
397 tab = astropy.table.Table({
398 "ra": [150.0], "diaSourceId": [100], "myExtra": [1]})
399 fake = _FakeTapService(tab)
400 with self.assertRaisesRegex(RuntimeError, "not in the SDM schema"):
401 PpdbTap(service=fake).load_sources(diaObjectId=42)
403 def test_explicit_columns_order_preserved(self):
404 # an explicit columns list is respected, not reordered to SDM order.
405 tab = astropy.table.Table({"dec": [1.0], "ra": [150.0]})
406 fake = _FakeTapService(tab)
407 result = PpdbTap(service=fake).load_sources(diaObjectId=42,
408 columns=["dec", "ra"])
409 self.assertEqual(result.colnames, ["dec", "ra"])
411 def test_objects_columns_reordered_to_sdm(self):
412 scrambled = astropy.table.Table({
413 "ra": [150.0], "diaObjectId": [1], "dec": [2.0],
414 "validityEndMjdTai": [60000.0]})
415 fake = _FakeTapService(scrambled)
416 result = PpdbTap(service=fake).load_objects()
417 self.assertEqual(
418 result.colnames,
419 ["diaObjectId", "validityEndMjdTai", "ra", "dec"])
421 def test_load_object_row_in_sdm_order(self):
422 scrambled = astropy.table.Table({
423 "ra": [150.0], "diaObjectId": [7], "dec": [2.0]})
424 fake = _FakeTapService(scrambled)
425 row = PpdbTap(service=fake).load_object(7)
426 self.assertEqual(row.colnames, ["diaObjectId", "ra", "dec"])
428 def test_id_list_path_reordered_to_sdm(self):
429 # the chunked many-ids path reorders per page before stacking.
430 scrambled = astropy.table.Table({
431 "band": ["r"], "diaObjectId": [42], "diaSourceId": [100],
432 "midpointMjdTai": [60000.0]})
433 fake = _FakeTapService(scrambled)
434 result = PpdbTap(service=fake).load_sources(diaObjectId=[1, 2],
435 id_chunk_size=1)
436 self.assertEqual(
437 result.colnames,
438 ["diaSourceId", "diaObjectId", "midpointMjdTai", "band"])
441class TestLightCurve(lsst.utils.tests.TestCase):
442 def test_assembles_object_and_series(self):
443 tables = [
444 _objects_table((42,)),
445 _sources_table(n=2),
446 _sources_table(n=3, id_column="diaForcedSourceId"),
447 ]
448 fake = _FakeTapService(tables)
449 light_curve = PpdbTap(service=fake).load_light_curve(42)
450 self.assertIsInstance(light_curve, DiaObjectLightCurve)
451 self.assertEqual(light_curve.diaObjectId, 42)
452 self.assertEqual(light_curve.n_sources, 2)
453 self.assertEqual(light_curve.n_forced_sources, 3)
454 # one query each for object, sources, forced sources.
455 self.assertEqual(len(fake.queries), 3)
456 self.assertIn("FROM ppdb.DiaObject", fake.queries[0])
459class TestMemory(lsst.utils.tests.MemoryTestCase):
460 pass
463def setup_module(module):
464 lsst.utils.tests.init()
467if __name__ == "__main__": 467 ↛ 468line 467 didn't jump to line 468 because the condition on line 467 was never true
468 lsst.utils.tests.init()
469 unittest.main()