Coverage for python / lsst / obs / base / ingest_tests.py: 26%

225 statements  

« prev     ^ index     » next       coverage.py v7.13.2, created at 2026-02-04 02:52 -0800

1# This file is part of obs_base. 

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 

22"""Base class for writing Gen3 raw data ingest tests.""" 

23 

24__all__ = ("IngestTestBase",) 

25 

26import abc 

27import os 

28import shutil 

29import tempfile 

30import unittest 

31 

32import lsst.afw.cameraGeom 

33import lsst.afw.cameraGeom.testUtils # For assertDetectorsEqual 

34import lsst.obs.base 

35from lsst.daf.butler import Butler, Registry 

36from lsst.daf.butler.cli.butler import cli as butlerCli 

37from lsst.daf.butler.cli.utils import LogCliRunner 

38from lsst.pipe.base import Instrument 

39from lsst.resources import ResourcePath 

40from lsst.utils import doImportType 

41 

42from . import script 

43 

44 

45class IngestTestBase(metaclass=abc.ABCMeta): 

46 """Base class for tests of gen3 ingest. Subclass from this, then 

47 `unittest.TestCase` to get a working test suite. 

48 """ 

49 

50 ingestDir = "" 

51 """Root path to ingest files into. Typically `obs_package/tests/`; the 

52 actual directory will be a tempdir under this one. 

53 """ 

54 

55 ingestDatasetTypeName = "raw" 

56 """The DatasetType to use for the ingest. 

57 

58 If this is not an Exposure dataset type the tests will be more limited. 

59 """ 

60 

61 dataIds = [] 

62 """list of butler data IDs of files that should have been ingested.""" 

63 

64 file = "" 

65 """Full path to a file to ingest in tests.""" 

66 

67 filterLabel = None 

68 """The lsst.afw.image.FilterLabel that should be returned by the above 

69 file.""" 

70 

71 rawIngestTask = "lsst.obs.base.RawIngestTask" 

72 """The task to use in the Ingest test.""" 

73 

74 curatedCalibrationDatasetTypes = None 

75 """List or tuple of Datasets types that should be present after calling 

76 writeCuratedCalibrations. If `None` writeCuratedCalibrations will 

77 not be called and the test will be skipped.""" 

78 

79 defineVisitsTask = lsst.obs.base.DefineVisitsTask 

80 """The task to use to define visits from groups of exposures. 

81 This is ignored if ``visits`` is `None`. 

82 """ 

83 

84 visits = {} 

85 """A dictionary mapping visit data IDs the lists of exposure data IDs that 

86 are associated with them. 

87 If this is empty (but not `None`), visit definition will be run but no 

88 visits will be expected (e.g. because no exposures are on-sky 

89 observations). 

90 """ 

91 

92 seed_config = None 

93 """Location of a seed configuration file to pass to butler create. 

94 

95 Useful if additional formatters or storage classes need to be defined. 

96 """ 

97 

98 @property 

99 @abc.abstractmethod 

100 def instrumentClassName(self): 

101 """The fully qualified instrument class name. 

102 

103 Returns 

104 ------- 

105 `str` 

106 The fully qualified instrument class name. 

107 """ 

108 pass 

109 

110 @property 

111 def instrumentClass(self): 

112 """The instrument class.""" 

113 return doImportType(self.instrumentClassName) 

114 

115 @property 

116 def instrumentName(self): 

117 """The name of the instrument. 

118 

119 Returns 

120 ------- 

121 `str` 

122 The name of the instrument. 

123 """ 

124 return self.instrumentClass.getName() 

125 

126 @classmethod 

127 def setUpClass(cls): 

128 # Use a temporary working directory. 

129 cls.root = tempfile.mkdtemp(dir=cls.ingestDir) 

130 cls._createRepo() 

131 

132 # Register the instrument and its static metadata. 

133 cls._registerInstrument() 

134 

135 # Determine the relevant datastore root to use for testing. 

136 with Butler.from_config(cls.root) as butler: 

137 roots = butler.get_datastore_roots() 

138 assert len(roots) == 1 # Only one datastore. 

139 cls.datastore_root = list(roots.values())[0] 

140 

141 def setUp(self): 

142 # Want a unique run name per test. 

143 self.outputRun = "raw_ingest_" + self.id() 

144 

145 @classmethod 

146 def tearDownClass(cls): 

147 if os.path.exists(cls.root): 

148 shutil.rmtree(cls.root, ignore_errors=True) 

149 

150 def verifyIngest(self, files=None, cli=False, fullCheck=False): 

151 """ 

152 Test that RawIngestTask ingested the expected files. 

153 

154 Parameters 

155 ---------- 

156 files : `list` [`str`], or None 

157 List of files to be ingested, or None to use ``self.file`` 

158 fullCheck : `bool`, optional 

159 If `True`, read the full raw dataset and check component 

160 consistency. If `False` check that a component can be read 

161 but do not read the entire raw exposure. 

162 

163 Notes 

164 ----- 

165 Reading all the ingested test data can be expensive. The code paths 

166 for reading the second raw are the same as reading the first so 

167 we do not gain anything by doing full checks of everything. 

168 Only read full pixel data for first dataset from file. 

169 Don't even do that if we are requested not to by the caller. 

170 This only really affects files that contain multiple datasets. 

171 """ 

172 butler = Butler.from_config(self.root, run=self.outputRun) 

173 self.enterContext(butler) 

174 datasets = list(butler.registry.queryDatasets(self.ingestDatasetTypeName, collections=self.outputRun)) 

175 self.assertEqual(len(datasets), len(self.dataIds)) 

176 

177 # Can check that the timespan in the day_obs matches the exposure 

178 # record. 

179 if "day_obs" in butler.dimensions: 

180 days = { 

181 (rec.instrument, rec.id): rec.timespan 

182 for rec in butler.registry.queryDimensionRecords("day_obs") 

183 } 

184 

185 exp_records = list(butler.registry.queryDimensionRecords("exposure")) 

186 for exp in exp_records: 

187 day_span = days[exp.instrument, exp.day_obs] 

188 if day_span is not None: 

189 self.assertTrue( 

190 day_span.contains(exp.timespan.begin), f"Timespan mismatch of {exp} and {day_span}" 

191 ) 

192 

193 # Get the URI to the first dataset and check it is inside the 

194 # datastore. 

195 datasetUri = butler.getURI(datasets[0]) 

196 self.assertIsNotNone(datasetUri.relative_to(self.datastore_root)) 

197 

198 # Get the relevant dataset type. 

199 datasetType = butler.get_dataset_type(self.ingestDatasetTypeName) 

200 

201 for dataId in self.dataIds: 

202 # For testing we only read the entire dataset the first time 

203 # round if this is an Exposure. If it's not an Exposure 

204 # we always read it completely but we don't read components 

205 # because for an arbitrary dataset type we can't easily tell 

206 # what component to test. 

207 

208 if not datasetType.storageClass.name.startswith("Exposure"): 

209 exposure = butler.get(self.ingestDatasetTypeName, dataId) 

210 # Could be anything so nothing to test by default 

211 continue 

212 

213 # Check that we can read metadata from a raw. 

214 metadata = butler.get(f"{self.ingestDatasetTypeName}.metadata", dataId) 

215 if not fullCheck: 

216 continue 

217 fullCheck = False 

218 exposure = butler.get(self.ingestDatasetTypeName, dataId) 

219 

220 # Comparing headers will not work directly because of header 

221 # fix up provenance. 

222 metadata_headers = metadata.toDict() 

223 exposure_headers = exposure.getMetadata().toDict() 

224 metadata_headers.pop("HIERARCH ASTRO METADATA FIX DATE", None) 

225 exposure_headers.pop("HIERARCH ASTRO METADATA FIX DATE", None) 

226 self.assertEqual(metadata_headers, exposure_headers) 

227 

228 # Since components follow a different code path we check that 

229 # WCS match and also we check that at least the shape 

230 # of the image is the same (rather than doing per-pixel equality) 

231 wcs = butler.get(f"{self.ingestDatasetTypeName}.wcs", dataId) 

232 self.assertEqual(wcs, exposure.getWcs()) 

233 

234 rawImage = butler.get(f"{self.ingestDatasetTypeName}.image", dataId) 

235 self.assertEqual(rawImage.getBBox(), exposure.getBBox()) 

236 

237 # Check that the filter label got the correct band. 

238 filterLabel = butler.get(f"{self.ingestDatasetTypeName}.filter", dataId) 

239 self.assertEqual(filterLabel, self.filterLabel) 

240 

241 # Check that the exposure's Detector is the same as the component 

242 # we would read (this is tricky for LSST, which modifies its 

243 # detector at read time; for most other cameras it should be 

244 # trivially satisfied. 

245 detector = butler.get(f"{self.ingestDatasetTypeName}.detector", dataId) 

246 self.assertDetectorsEqual(detector, exposure.getDetector(), compareTransforms=False) 

247 

248 self.checkRepo(files=files) 

249 

250 def checkRepo(self, files=None): 

251 """Check the state of the repository after ingest. 

252 

253 This is an optional hook provided for subclasses; by default it does 

254 nothing. 

255 

256 Parameters 

257 ---------- 

258 files : `list` [`str`], or None 

259 List of files to be ingested, or None to use ``self.file`` 

260 """ 

261 return 

262 

263 @classmethod 

264 def _createRepo(cls): 

265 """Use the Click `testing` module to call the butler command line api 

266 to create a repository. 

267 """ 

268 runner = LogCliRunner() 

269 args = [] 

270 if cls.seed_config: 

271 args.extend(["--seed-config", cls.seed_config]) 

272 result = runner.invoke(butlerCli, ["create", cls.root, *args]) 

273 # Classmethod so assertEqual does not work. 

274 assert result.exit_code == 0, f"output: {result.output} exception: {result.exception}" 

275 

276 def _ingestRaws(self, transfer, file=None, skip_existing=True): 

277 """Use the Click `testing` module to call the butler command line api 

278 to ingest raws. 

279 

280 Parameters 

281 ---------- 

282 transfer : `str` 

283 The external data transfer type. 

284 file : `str` 

285 Path to a file to ingest instead of the default associated with 

286 the object. 

287 """ 

288 if file is None: 

289 file = self.file 

290 

291 args = [ 

292 "ingest-raws", 

293 self.root, 

294 file, 

295 "--output-run", 

296 self.outputRun, 

297 "--transfer", 

298 transfer, 

299 "--ingest-task", 

300 self.rawIngestTask, 

301 ] 

302 if not skip_existing: 

303 args.append("--no-skip-existing") 

304 runner = LogCliRunner() 

305 result = runner.invoke(butlerCli, args) 

306 self.assertEqual(result.exit_code, 0, f"output: {result.output} exception: {result.exception}") 

307 

308 @classmethod 

309 def _registerInstrument(cls): 

310 """Use the Click `testing` module to call the butler command line api 

311 to register the instrument. 

312 """ 

313 runner = LogCliRunner() 

314 result = runner.invoke(butlerCli, ["register-instrument", cls.root, cls.instrumentClassName]) 

315 # Classmethod so assertEqual does not work. 

316 assert result.exit_code == 0, f"output: {result.output} exception: {result.exception}" 

317 

318 def _writeCuratedCalibrations(self): 

319 """Use the Click `testing` module to call the butler command line api 

320 to write curated calibrations. 

321 """ 

322 runner = LogCliRunner() 

323 result = runner.invoke( 

324 butlerCli, ["write-curated-calibrations", self.root, self.instrumentName, "test"] 

325 ) 

326 self.assertEqual(result.exit_code, 0, f"output: {result.output} exception: {result.exception}") 

327 

328 def testLink(self): 

329 self._ingestRaws(transfer="link") 

330 self.verifyIngest() 

331 

332 def testSymLink(self): 

333 self._ingestRaws(transfer="symlink") 

334 self.verifyIngest() 

335 

336 def testDirect(self): 

337 self._ingestRaws(transfer="direct") 

338 

339 # Check that it really did have a URI outside of datastore. 

340 srcUri = ResourcePath(self.file, forceAbsolute=True) 

341 butler = Butler.from_config(self.root, run=self.outputRun) 

342 self.enterContext(butler) 

343 datasets = list(butler.registry.queryDatasets(self.ingestDatasetTypeName, collections=self.outputRun)) 

344 datastoreUri = butler.getURI(datasets[0]) 

345 self.assertEqual(datastoreUri, srcUri) 

346 

347 def testCopy(self): 

348 self._ingestRaws(transfer="copy") 

349 # Only test full read of raws for the copy test. No need to do it 

350 # in the other tests since the formatter will be the same in all 

351 # cases. 

352 self.verifyIngest(fullCheck=True) 

353 

354 def testHardLink(self): 

355 try: 

356 self._ingestRaws(transfer="hardlink") 

357 # Running ingest through the Click testing infrastructure causes 

358 # the original exception indicating that we can't hard-link 

359 # on this filesystem to be turned into a nonzero exit code, which 

360 # then trips the test assertion. 

361 except (AssertionError, PermissionError) as err: 

362 raise unittest.SkipTest( 

363 "Skipping hard-link test because input data is on a different filesystem." 

364 ) from err 

365 self.verifyIngest() 

366 

367 def testInPlace(self): 

368 """Test that files already in the directory can be added to the 

369 registry in-place. 

370 """ 

371 butler = Butler.from_config(self.root, run=self.outputRun) 

372 self.enterContext(butler) 

373 

374 # If the test uses an index file the index file needs to also 

375 # appear in the datastore root along with the file to be ingested. 

376 # In that scenario the file name being used for ingest can not 

377 # be modified and must have the same name as found in the index 

378 # file itself. 

379 source_file_uri = ResourcePath(self.file) 

380 index_file = source_file_uri.dirname().join("_index.json") 

381 pathInStore = source_file_uri.basename() 

382 if index_file.exists(): 

383 os.symlink(index_file.ospath, self.datastore_root.join("_index.json").ospath) 

384 else: 

385 # No index file so we are free to pick any name. 

386 pathInStore = "prefix-" + pathInStore 

387 

388 # Create a symlink to the original file so that it looks like it 

389 # is now inside the datastore. 

390 newPath = self.datastore_root.join(pathInStore) 

391 os.symlink(os.path.abspath(self.file), newPath.ospath) 

392 

393 # If there is a sidecar file it needs to be linked in as well 

394 # since ingest code does not follow symlinks. 

395 sidecar_uri = ResourcePath(source_file_uri).updatedExtension(".json") 

396 if sidecar_uri.exists(): 

397 newSidecar = ResourcePath(newPath).updatedExtension(".json") 

398 os.symlink(sidecar_uri.ospath, newSidecar.ospath) 

399 

400 # Run ingest with auto mode since that should automatically determine 

401 # that an in-place ingest is happening. 

402 self._ingestRaws(transfer="auto", file=newPath.ospath) 

403 self.verifyIngest() 

404 

405 # Recreate a butler post-ingest (the earlier one won't see the 

406 # ingested files). 

407 butler = Butler.from_config(self.root, run=self.outputRun) 

408 self.enterContext(butler) 

409 

410 # Check that the URI associated with this path is the right one. 

411 uri = butler.getURI(self.ingestDatasetTypeName, self.dataIds[0]) 

412 self.assertEqual(uri.relative_to(self.datastore_root), pathInStore) 

413 

414 def testFailOnConflict(self): 

415 """Re-ingesting the same data into the repository should fail.""" 

416 self._ingestRaws(transfer="symlink") 

417 self._ingestRaws(transfer="symlink") # Default is to skip. 

418 with self.assertRaises(AssertionError): 

419 self._ingestRaws(transfer="symlink", skip_existing=False) 

420 

421 def testWriteCuratedCalibrations(self): 

422 """Test that we can ingest the curated calibrations, and read them 

423 with `loadCamera` both before and after. 

424 """ 

425 if self.curatedCalibrationDatasetTypes is None: 

426 raise unittest.SkipTest("Class requests disabling of writeCuratedCalibrations test") 

427 

428 butler = Butler.from_config(self.root, writeable=False) 

429 self.enterContext(butler) 

430 collection = self.instrumentClass().makeCalibrationCollectionName("test") 

431 

432 # Trying to load a camera with a data ID not known to the registry 

433 # is an error, because we can't get any temporal information. 

434 with self.assertRaises(LookupError): 

435 lsst.obs.base.loadCamera(butler, {"exposure": 0}, collections=collection) 

436 

437 # Ingest raws in order to get some exposure records. 

438 self._ingestRaws(transfer="auto") 

439 

440 # Load camera should returned an unversioned camera because there's 

441 # nothing in the repo. 

442 camera, isVersioned = lsst.obs.base.loadCamera(butler, self.dataIds[0], collections=collection) 

443 self.assertFalse(isVersioned) 

444 self.assertIsInstance(camera, lsst.afw.cameraGeom.Camera) 

445 

446 self._writeCuratedCalibrations() 

447 

448 # Make a new butler instance to make sure we don't have any stale 

449 # caches (e.g. of DatasetTypes). Note that we didn't give 

450 # _writeCuratedCalibrations the butler instance we had, because it's 

451 # trying to test the CLI interface anyway. 

452 butler = Butler.from_config(self.root, writeable=False) 

453 self.enterContext(butler) 

454 

455 instrumentClass = self.instrumentClass() 

456 calibration_names = instrumentClass.getCuratedCalibrationNames() 

457 

458 for datasetTypeName in self.curatedCalibrationDatasetTypes: 

459 with self.subTest(dtype=datasetTypeName): 

460 found = list( 

461 butler.registry.queryDatasetAssociations( 

462 datasetTypeName, 

463 collections=collection, 

464 ) 

465 ) 

466 self.assertGreater(len(found), 0, f"Checking {datasetTypeName}") 

467 self.assertIn(datasetTypeName, calibration_names) 

468 

469 # Load camera should returned the versioned camera from the repo. 

470 camera, isVersioned = lsst.obs.base.loadCamera(butler, self.dataIds[0], collections=collection) 

471 self.assertTrue(isVersioned) 

472 self.assertIsInstance(camera, lsst.afw.cameraGeom.Camera) 

473 

474 def testDefineVisits(self): 

475 if self.visits is None: 

476 self.skipTest("Expected visits were not defined.") 

477 self._ingestRaws(transfer="link") 

478 

479 # Check that obscore table (if configured) has correct contents. 

480 butler = Butler.from_config(self.root, run=self.outputRun) 

481 self.enterContext(butler) 

482 self._check_obscore(butler.registry, has_visits=False) 

483 

484 # Calling defineVisits tests the implementation of the butler command 

485 # line interface "define-visits" subcommand. Functions in the script 

486 # folder are generally considered protected and should not be used 

487 # as public api. 

488 script.defineVisits( 

489 self.root, 

490 config_file=None, 

491 collections=self.outputRun, 

492 instrument=self.instrumentName, 

493 ) 

494 

495 # Test that we got the visits we expected. 

496 visits = butler.registry.queryDataIds(["visit"]).expanded().toSet() 

497 self.assertCountEqual(visits, self.visits.keys()) 

498 instr = Instrument.from_string(self.instrumentName, butler.registry) 

499 camera = instr.getCamera() 

500 for foundVisit, (expectedVisit, expectedExposures) in zip(visits, self.visits.items(), strict=True): 

501 # Test that this visit is associated with the expected exposures. 

502 foundExposures = ( 

503 butler.registry.queryDataIds(["exposure"], dataId=expectedVisit).expanded().toSet() 

504 ) 

505 self.assertCountEqual(foundExposures, expectedExposures) 

506 # Test that we have a visit region, and that it contains all of the 

507 # detector+visit regions. 

508 self.assertIsNotNone(foundVisit.region) 

509 detectorVisitDataIds = ( 

510 butler.registry.queryDataIds(["visit", "detector"], dataId=expectedVisit).expanded().toSet() 

511 ) 

512 self.assertEqual(len(detectorVisitDataIds), len(camera)) 

513 for dataId in detectorVisitDataIds: 

514 self.assertTrue(foundVisit.region.contains(dataId.region)) 

515 

516 # Check obscore table again. 

517 self._check_obscore(butler.registry, has_visits=True) 

518 

519 def _check_obscore(self, registry: Registry, has_visits: bool) -> None: 

520 """Verify contents of obscore table.""" 

521 return