Coverage for tests/test_ppdb.py: 99%
300 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:33 +0000
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.assertNotIn("WHERE", query)
159 self.assertIn("TOP 100000", query)
160 self.assertIn("ORDER BY diaObjectId", query)
162 def test_cone_search(self):
163 fake = _FakeTapService(_objects_table())
164 PpdbTap(service=fake).load_objects(ra=150.0, dec=2.5, radius=1.0)
165 query = fake.queries[-1]
166 self.assertIn("CONTAINS(POINT('ICRS', ra, dec)", query)
167 self.assertIn("CIRCLE('ICRS', 150.0, 2.5, 1.0)", query)
169 def test_partial_cone_raises(self):
170 fake = _FakeTapService(_objects_table())
171 with self.assertRaisesRegex(ValueError, "ra, dec, and radius"):
172 PpdbTap(service=fake).load_objects(ra=150.0, dec=2.5)
174 def test_exposure_region(self):
175 fake = _FakeTapService(_objects_table())
176 PpdbTap(service=fake).load_objects(exposure=_make_exposure())
177 self.assertIn("CONTAINS(POINT('ICRS', ra, dec)", fake.queries[-1])
179 def test_exposure_and_cone_conflict(self):
180 fake = _FakeTapService(_objects_table())
181 with self.assertRaisesRegex(ValueError, "either exposure="):
182 PpdbTap(service=fake).load_objects(exposure=_make_exposure(),
183 ra=150.0, dec=2.5, radius=1.0)
185 def test_no_limit_omits_top(self):
186 fake = _FakeTapService(_objects_table())
187 PpdbTap(service=fake).load_objects(limit=None)
188 self.assertNotIn("TOP", fake.queries[-1])
190 def test_duplicate_current_versions_warns(self):
191 fake = _FakeTapService(_objects_table((1, 1, 2)))
192 ppdb = PpdbTap(service=fake)
193 with self.assertLogs("lsst.analysis.ap.ppdb", level="WARNING") as cm:
194 ppdb.load_objects()
195 self.assertTrue(any("duplicate current versions" in m
196 for m in cm.output))
198 def test_truncation_warns(self):
199 fake = _FakeTapService(_objects_table((1, 2)))
200 ppdb = PpdbTap(service=fake)
201 with self.assertLogs("lsst.analysis.ap.ppdb", level="WARNING") as cm:
202 ppdb.load_objects(limit=2)
203 self.assertTrue(any("row limit" in m for m in cm.output))
205 def test_load_object_missing_raises(self):
206 fake = _FakeTapService(astropy.table.Table())
207 with self.assertRaisesRegex(RuntimeError, "diaObjectId=999 not found"):
208 PpdbTap(service=fake).load_object(999)
210 def test_load_object_returns_row(self):
211 fake = _FakeTapService(_objects_table((7,)))
212 row = PpdbTap(service=fake).load_object(7)
213 self.assertEqual(row["diaObjectId"], 7)
214 self.assertIn("WHERE diaObjectId = 7", fake.queries[-1])
217class TestLoadSources(lsst.utils.tests.TestCase):
218 def test_sources_for_object(self):
219 fake = _FakeTapService(_sources_table())
220 PpdbTap(service=fake).load_sources_for_object(42)
221 query = fake.queries[-1]
222 self.assertIn("FROM ppdb.DiaSource", query)
223 self.assertIn("diaObjectId = 42", query)
224 self.assertIn("ORDER BY midpointMjdTai", query)
226 def test_forced_sources_for_object(self):
227 fake = _FakeTapService(_sources_table(id_column="diaForcedSourceId"))
228 PpdbTap(service=fake).load_forced_sources_for_object(42)
229 self.assertIn("FROM ppdb.DiaForcedSource", fake.queries[-1])
231 def test_band_filter(self):
232 fake = _FakeTapService(_sources_table())
233 PpdbTap(service=fake).load_sources(diaObjectId=42, bands=["r", "g"])
234 query = fake.queries[-1]
235 self.assertIn("diaObjectId = 42", query)
236 self.assertIn("band IN ('r', 'g')", query)
238 def test_invalid_band_raises(self):
239 fake = _FakeTapService(_sources_table())
240 with self.assertRaisesRegex(ValueError, "Unknown band"):
241 PpdbTap(service=fake).load_sources(diaObjectId=42, bands=["x"])
243 def test_time_window(self):
244 fake = _FakeTapService(_sources_table())
245 PpdbTap(service=fake).load_sources(diaObjectId=42, mjd_begin=100.0,
246 mjd_end=200.0)
247 query = fake.queries[-1]
248 self.assertIn("midpointMjdTai >= 100.0", query)
249 self.assertIn("midpointMjdTai < 200.0", query)
251 def test_object_id_required(self):
252 fake = _FakeTapService(_sources_table())
253 # omitting it entirely is a TypeError (required keyword)
254 with self.assertRaises(TypeError):
255 PpdbTap(service=fake).load_sources()
256 # passing None explicitly gives a guiding ValueError
257 with self.assertRaisesRegex(ValueError, "diaObjectId is required"):
258 PpdbTap(service=fake).load_sources(diaObjectId=None)
260 def test_load_source_missing_raises(self):
261 fake = _FakeTapService(astropy.table.Table())
262 with self.assertRaisesRegex(RuntimeError, "diaSourceId=5 not found"):
263 PpdbTap(service=fake).load_source(5)
265 def test_load_source_returns_row(self):
266 fake = _FakeTapService(_sources_table(n=1))
267 row = PpdbTap(service=fake).load_source(100)
268 self.assertEqual(row["diaSourceId"], 100)
270 def test_id_list_is_chunked(self):
271 fake = _FakeTapService(_sources_table(n=2))
272 result = PpdbTap(service=fake).load_sources(
273 diaObjectId=[1, 2, 3, 4, 5], id_chunk_size=2)
274 # 5 ids in chunks of 2 -> 3 IN queries, results concatenated.
275 self.assertEqual(len(fake.queries), 3)
276 self.assertIn("diaObjectId IN (1, 2)", fake.queries[0])
277 self.assertIn("diaObjectId IN (3, 4)", fake.queries[1])
278 self.assertIn("diaObjectId IN (5)", fake.queries[2])
279 self.assertEqual(len(result), 6)
281 def test_empty_id_list_returns_empty(self):
282 fake = _FakeTapService(_sources_table(n=0))
283 result = PpdbTap(service=fake).load_sources(diaObjectId=[])
284 self.assertEqual(len(result), 0)
285 self.assertIn("1 = 0", fake.queries[-1])
288class TestRegionAndCone(lsst.utils.tests.TestCase):
289 def test_sources_for_region_two_step(self):
290 fake = _FakeTapService([_objects_table((10, 20, 30)),
291 _sources_table(n=2)])
292 PpdbTap(service=fake).load_sources_for_region(ra=150.0, dec=2.5,
293 radius=0.1)
294 # first: cone on DiaObject (the only table allowing spatial search).
295 obj_query = fake.queries[0]
296 self.assertIn("FROM ppdb.DiaObject", obj_query)
297 self.assertIn("CONTAINS(POINT('ICRS', ra, dec)", obj_query)
298 # second: sources by id, with no spatial predicate.
299 src_query = fake.queries[1]
300 self.assertIn("FROM ppdb.DiaSource", src_query)
301 self.assertIn("diaObjectId IN (10, 20, 30)", src_query)
302 self.assertNotIn("CONTAINS", src_query)
304 def test_forced_sources_for_region_with_exposure(self):
305 fake = _FakeTapService([
306 _objects_table((7,)),
307 _sources_table(n=1, id_column="diaForcedSourceId"),
308 ])
309 PpdbTap(service=fake).load_forced_sources_for_region(
310 exposure=_make_exposure())
311 self.assertIn("FROM ppdb.DiaObject", fake.queries[0])
312 self.assertIn("CONTAINS", fake.queries[0])
313 self.assertIn("FROM ppdb.DiaForcedSource", fake.queries[1])
314 self.assertIn("diaObjectId IN (7)", fake.queries[1])
316 def test_region_with_no_objects(self):
317 fake = _FakeTapService([astropy.table.Table(), _sources_table(n=0)])
318 result = PpdbTap(service=fake).load_sources_for_region(
319 ra=1.0, dec=1.0, radius=0.01)
320 self.assertEqual(len(result), 0)
321 self.assertIn("1 = 0", fake.queries[1])
323 def test_sources_by_cone_warns_and_queries_source_table(self):
324 fake = _FakeTapService(_sources_table(n=2))
325 ppdb = PpdbTap(service=fake)
326 with self.assertLogs("lsst.analysis.ap.ppdb", level="WARNING") as cm:
327 ppdb.load_sources_by_cone(ra=150.0, dec=2.5, radius=0.05)
328 self.assertTrue(any("prototype-only" in m for m in cm.output))
329 query = fake.queries[-1]
330 self.assertIn("FROM ppdb.DiaSource", query)
331 self.assertIn("CONTAINS(POINT('ICRS', ra, dec)", query)
333 def test_forced_sources_by_cone(self):
334 fake = _FakeTapService(
335 _sources_table(n=1, id_column="diaForcedSourceId"))
336 ppdb = PpdbTap(service=fake)
337 with self.assertLogs("lsst.analysis.ap.ppdb", level="WARNING"):
338 ppdb.load_forced_sources_by_cone(exposure=_make_exposure())
339 self.assertIn("FROM ppdb.DiaForcedSource", fake.queries[-1])
340 self.assertIn("CONTAINS", fake.queries[-1])
342 def test_by_cone_requires_a_cone(self):
343 fake = _FakeTapService(_sources_table())
344 with self.assertRaisesRegex(ValueError, "cone search requires"):
345 PpdbTap(service=fake).load_sources_by_cone()
347 def test_sources_by_time_window_warns_and_queries_source_table(self):
348 fake = _FakeTapService(_sources_table(n=2))
349 ppdb = PpdbTap(service=fake)
350 with self.assertLogs("lsst.analysis.ap.ppdb", level="WARNING") as cm:
351 ppdb.load_sources_by_time_window(mjd_begin=60000.0, mjd_end=60001.0)
352 self.assertTrue(any("prototype-only" in m for m in cm.output))
353 query = fake.queries[-1]
354 self.assertIn("FROM ppdb.DiaSource", query)
355 self.assertIn("midpointMjdTai >= 60000.0", query)
356 self.assertIn("midpointMjdTai < 60001.0", query)
357 self.assertNotIn("diaObjectId", query)
358 self.assertNotIn("CONTAINS", query)
360 def test_forced_sources_by_time_window_open_ended(self):
361 fake = _FakeTapService(
362 _sources_table(n=1, id_column="diaForcedSourceId"))
363 ppdb = PpdbTap(service=fake)
364 with self.assertLogs("lsst.analysis.ap.ppdb", level="WARNING"):
365 ppdb.load_forced_sources_by_time_window(mjd_begin=60000.0)
366 query = fake.queries[-1]
367 self.assertIn("FROM ppdb.DiaForcedSource", query)
368 self.assertIn("midpointMjdTai >= 60000.0", query)
370 def test_by_time_window_requires_a_bound(self):
371 fake = _FakeTapService(_sources_table())
372 with self.assertRaisesRegex(ValueError, "time-window search requires"):
373 PpdbTap(service=fake).load_sources_by_time_window()
376class TestColumnOrdering(lsst.utils.tests.TestCase):
377 def test_sources_columns_reordered_to_sdm(self):
378 # alphabetical input columns -> SDM-schema order out.
379 scrambled = astropy.table.Table({
380 "band": ["r"], "dec": [1.0], "diaObjectId": [42],
381 "diaSourceId": [100], "midpointMjdTai": [60000.0], "ra": [150.0]})
382 fake = _FakeTapService(scrambled)
383 result = PpdbTap(service=fake).load_sources(diaObjectId=42)
384 self.assertEqual(result.colnames,
385 ["diaSourceId", "diaObjectId", "midpointMjdTai",
386 "ra", "dec", "band"])
388 def test_unknown_columns_raise(self):
389 # a column not in the SDM schema signals schema drift -> error.
390 tab = astropy.table.Table({
391 "ra": [150.0], "diaSourceId": [100], "myExtra": [1]})
392 fake = _FakeTapService(tab)
393 with self.assertRaisesRegex(RuntimeError, "not in the SDM schema"):
394 PpdbTap(service=fake).load_sources(diaObjectId=42)
396 def test_explicit_columns_order_preserved(self):
397 # an explicit columns list is respected, not reordered to SDM order.
398 tab = astropy.table.Table({"dec": [1.0], "ra": [150.0]})
399 fake = _FakeTapService(tab)
400 result = PpdbTap(service=fake).load_sources(diaObjectId=42,
401 columns=["dec", "ra"])
402 self.assertEqual(result.colnames, ["dec", "ra"])
404 def test_objects_columns_reordered_to_sdm(self):
405 scrambled = astropy.table.Table({
406 "ra": [150.0], "diaObjectId": [1], "dec": [2.0],
407 "validityStartMjdTai": [60000.0]})
408 fake = _FakeTapService(scrambled)
409 result = PpdbTap(service=fake).load_objects()
410 self.assertEqual(
411 result.colnames,
412 ["diaObjectId", "validityStartMjdTai", "ra", "dec"])
414 def test_load_object_row_in_sdm_order(self):
415 scrambled = astropy.table.Table({
416 "ra": [150.0], "diaObjectId": [7], "dec": [2.0]})
417 fake = _FakeTapService(scrambled)
418 row = PpdbTap(service=fake).load_object(7)
419 self.assertEqual(row.colnames, ["diaObjectId", "ra", "dec"])
421 def test_id_list_path_reordered_to_sdm(self):
422 # the chunked many-ids path reorders per page before stacking.
423 scrambled = astropy.table.Table({
424 "band": ["r"], "diaObjectId": [42], "diaSourceId": [100],
425 "midpointMjdTai": [60000.0]})
426 fake = _FakeTapService(scrambled)
427 result = PpdbTap(service=fake).load_sources(diaObjectId=[1, 2],
428 id_chunk_size=1)
429 self.assertEqual(
430 result.colnames,
431 ["diaSourceId", "diaObjectId", "midpointMjdTai", "band"])
434class TestLightCurve(lsst.utils.tests.TestCase):
435 def test_assembles_object_and_series(self):
436 tables = [
437 _objects_table((42,)),
438 _sources_table(n=2),
439 _sources_table(n=3, id_column="diaForcedSourceId"),
440 ]
441 fake = _FakeTapService(tables)
442 light_curve = PpdbTap(service=fake).load_light_curve(42)
443 self.assertIsInstance(light_curve, DiaObjectLightCurve)
444 self.assertEqual(light_curve.diaObjectId, 42)
445 self.assertEqual(light_curve.n_sources, 2)
446 self.assertEqual(light_curve.n_forced_sources, 3)
447 # one query each for object, sources, forced sources.
448 self.assertEqual(len(fake.queries), 3)
449 self.assertIn("FROM ppdb.DiaObject", fake.queries[0])
452class TestMemory(lsst.utils.tests.MemoryTestCase):
453 pass
456def setup_module(module):
457 lsst.utils.tests.init()
460if __name__ == "__main__": 460 ↛ 461line 460 didn't jump to line 461 because the condition on line 460 was never true
461 lsst.utils.tests.init()
462 unittest.main()