Coverage for tests/test_simpleButler.py: 12%

422 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-10-10 01:29 -0700

1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://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 software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <http://www.gnu.org/licenses/>. 

27 

28from __future__ import annotations 

29 

30import json 

31import os 

32import re 

33import tempfile 

34import unittest 

35import urllib.parse 

36from typing import Any 

37 

38try: 

39 import numpy as np 

40except ImportError: 

41 np = None 

42 

43import astropy.time 

44 

45from lsst.daf.butler import ( 

46 Butler, 

47 ButlerConfig, 

48 CollectionType, 

49 DataCoordinate, 

50 DatasetId, 

51 DatasetRef, 

52 DatasetType, 

53 LabeledButlerFactory, 

54 StorageClass, 

55 Timespan, 

56) 

57from lsst.daf.butler.datastore.file_templates import FileTemplate 

58from lsst.daf.butler.registry import RegistryConfig, RegistryDefaults, _RegistryFactory 

59from lsst.daf.butler.tests import DatastoreMock 

60from lsst.daf.butler.tests.utils import TestCaseMixin, makeTestTempDir, mock_env, removeTestTempDir 

61 

62try: 

63 from lsst.daf.butler.tests.server import create_test_server 

64except ImportError: 

65 create_test_server = None 

66 

67TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

68 

69 

70class SimpleButlerTests(TestCaseMixin): 

71 """Tests for butler (including import/export functionality) that should not 

72 depend on the Registry Database backend or Datastore implementation, and 

73 can instead utilize an in-memory SQLite Registry and a mocked Datastore. 

74 """ 

75 

76 datasetsImportFile = "datasets.yaml" 

77 

78 supportsCollectionRegex: bool = True 

79 """True if the registry class being tested supports regex searches for 

80 collections.""" 

81 

82 def makeButler(self, writeable: bool = False) -> Butler: 

83 raise NotImplementedError() 

84 

85 def comparableRef(self, ref: DatasetRef) -> DatasetRef: 

86 """Return a DatasetRef that can be compared to a DatasetRef from 

87 other repository. 

88 

89 For repositories that do not support round-trip of ID values this 

90 method returns unresolved DatasetRef, for round-trip-safe repos it 

91 returns unchanged ref. 

92 """ 

93 return ref 

94 

95 def testReadBackwardsCompatibility(self): 

96 """Test that we can read an export file written by a previous version 

97 and commit to the daf_butler git repo. 

98 

99 Notes 

100 ----- 

101 At present this export file includes only dimension data, not datasets, 

102 which greatly limits the usefulness of this test. We should address 

103 this at some point, but I think it's best to wait for the changes to 

104 the export format required for CALIBRATION collections to land. 

105 """ 

106 butler = self.makeButler(writeable=True) 

107 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "hsc-rc2-subset.yaml")) 

108 # Spot-check a few things, but the most important test is just that 

109 # the above does not raise. 

110 self.assertGreaterEqual( 

111 {record.id for record in butler.registry.queryDimensionRecords("detector", instrument="HSC")}, 

112 set(range(104)), # should have all science CCDs; may have some focus ones. 

113 ) 

114 self.assertGreaterEqual( 

115 { 

116 (record.id, record.physical_filter) 

117 for record in butler.registry.queryDimensionRecords("visit", instrument="HSC") 

118 }, 

119 { 

120 (27136, "HSC-Z"), 

121 (11694, "HSC-G"), 

122 (23910, "HSC-R"), 

123 (11720, "HSC-Y"), 

124 (23900, "HSC-R"), 

125 (22646, "HSC-Y"), 

126 (1248, "HSC-I"), 

127 (19680, "HSC-I"), 

128 (1240, "HSC-I"), 

129 (424, "HSC-Y"), 

130 (19658, "HSC-I"), 

131 (344, "HSC-Y"), 

132 (1218, "HSC-R"), 

133 (1190, "HSC-Z"), 

134 (23718, "HSC-R"), 

135 (11700, "HSC-G"), 

136 (26036, "HSC-G"), 

137 (23872, "HSC-R"), 

138 (1170, "HSC-Z"), 

139 (1876, "HSC-Y"), 

140 }, 

141 ) 

142 

143 def testDatasetTransfers(self): 

144 """Test exporting all datasets from a repo and then importing them all 

145 back in again. 

146 """ 

147 # Import data to play with. 

148 butler1 = self.makeButler(writeable=True) 

149 butler1.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml") 

150 butler1.import_(filename=f"resource://lsst.daf.butler/tests/registry_data/{self.datasetsImportFile}") 

151 with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml") as file: 

152 # Export all datasets. 

153 with butler1.export(filename=file.name) as exporter: 

154 exporter.saveDatasets(butler1.registry.queryDatasets(..., collections=...)) 

155 # Import it all again. 

156 butler2 = self.makeButler(writeable=True) 

157 butler2.import_(filename=file.name) 

158 datasets1 = list(butler1.registry.queryDatasets(..., collections=...)) 

159 datasets2 = list(butler2.registry.queryDatasets(..., collections=...)) 

160 self.assertTrue(all(isinstance(ref.id, DatasetId) for ref in datasets1)) 

161 self.assertTrue(all(isinstance(ref.id, DatasetId) for ref in datasets2)) 

162 self.assertCountEqual( 

163 [self.comparableRef(ref) for ref in datasets1], 

164 [self.comparableRef(ref) for ref in datasets2], 

165 ) 

166 

167 def testImportTwice(self): 

168 """Test exporting dimension records and datasets from a repo and then 

169 importing them all back in again twice. 

170 """ 

171 # Import data to play with. 

172 butler1 = self.makeButler(writeable=True) 

173 butler1.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml") 

174 butler1.import_(filename=f"resource://lsst.daf.butler/tests/registry_data/{self.datasetsImportFile}") 

175 with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as file: 

176 # Export all datasets. 

177 with butler1.export(filename=file.name) as exporter: 

178 exporter.saveDatasets(butler1.registry.queryDatasets(..., collections=...)) 

179 butler2 = self.makeButler(writeable=True) 

180 # Import it once. 

181 butler2.import_(filename=file.name) 

182 # Import it again 

183 butler2.import_(filename=file.name) 

184 datasets1 = list(butler1.registry.queryDatasets(..., collections=...)) 

185 datasets2 = list(butler2.registry.queryDatasets(..., collections=...)) 

186 self.assertTrue(all(isinstance(ref.id, DatasetId) for ref in datasets1)) 

187 self.assertTrue(all(isinstance(ref.id, DatasetId) for ref in datasets2)) 

188 self.assertCountEqual( 

189 [self.comparableRef(ref) for ref in datasets1], 

190 [self.comparableRef(ref) for ref in datasets2], 

191 ) 

192 

193 def testCollectionTransfers(self): 

194 """Test exporting and then importing collections of various types.""" 

195 # Populate a registry with some datasets. 

196 butler1 = self.makeButler(writeable=True) 

197 butler1.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml") 

198 butler1.import_(filename=f"resource://lsst.daf.butler/tests/registry_data/{self.datasetsImportFile}") 

199 registry1 = butler1.registry 

200 # Add some more collections. 

201 registry1.registerRun("run1") 

202 registry1.registerCollection("tag1", CollectionType.TAGGED) 

203 registry1.registerCollection("calibration1", CollectionType.CALIBRATION) 

204 registry1.registerCollection("chain1", CollectionType.CHAINED) 

205 registry1.registerCollection("chain2", CollectionType.CHAINED) 

206 registry1.setCollectionChain("chain1", ["tag1", "run1", "chain2"]) 

207 registry1.setCollectionChain("chain2", ["calibration1", "run1"]) 

208 # Associate some datasets into the TAGGED and CALIBRATION collections. 

209 flats1 = list(registry1.queryDatasets("flat", collections=...)) 

210 registry1.associate("tag1", flats1) 

211 t1 = astropy.time.Time("2020-01-01T01:00:00", format="isot", scale="tai") 

212 t2 = astropy.time.Time("2020-01-01T02:00:00", format="isot", scale="tai") 

213 t3 = astropy.time.Time("2020-01-01T03:00:00", format="isot", scale="tai") 

214 bias1a = registry1.findDataset("bias", instrument="Cam1", detector=1, collections="imported_g") 

215 bias2a = registry1.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g") 

216 bias3a = registry1.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g") 

217 bias2b = registry1.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r") 

218 bias3b = registry1.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r") 

219 registry1.certify("calibration1", [bias2a, bias3a], Timespan(t1, t2)) 

220 registry1.certify("calibration1", [bias2b], Timespan(t2, None)) 

221 registry1.certify("calibration1", [bias3b], Timespan(t2, t3)) 

222 registry1.certify("calibration1", [bias1a], Timespan.makeEmpty()) 

223 

224 with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml") as file: 

225 # Export all collections, and some datasets. 

226 with butler1.export(filename=file.name) as exporter: 

227 # Sort results to put chain1 before chain2, which is 

228 # intentionally not topological order. 

229 for collection in sorted(registry1.queryCollections()): 

230 exporter.saveCollection(collection) 

231 exporter.saveDatasets(flats1) 

232 exporter.saveDatasets([bias1a, bias2a, bias2b, bias3a, bias3b]) 

233 # Import them into a new registry. 

234 butler2 = self.makeButler(writeable=True) 

235 butler2.import_(filename=file.name) 

236 registry2 = butler2.registry 

237 # Check that it all round-tripped, starting with the collections 

238 # themselves. 

239 self.assertIs(registry2.getCollectionType("run1"), CollectionType.RUN) 

240 self.assertIs(registry2.getCollectionType("tag1"), CollectionType.TAGGED) 

241 self.assertIs(registry2.getCollectionType("calibration1"), CollectionType.CALIBRATION) 

242 self.assertIs(registry2.getCollectionType("chain1"), CollectionType.CHAINED) 

243 self.assertIs(registry2.getCollectionType("chain2"), CollectionType.CHAINED) 

244 self.assertEqual( 

245 list(registry2.getCollectionChain("chain1")), 

246 ["tag1", "run1", "chain2"], 

247 ) 

248 self.assertEqual( 

249 list(registry2.getCollectionChain("chain2")), 

250 ["calibration1", "run1"], 

251 ) 

252 # Check that tag collection contents are the same. 

253 self.maxDiff = None 

254 self.assertCountEqual( 

255 [self.comparableRef(ref) for ref in registry1.queryDatasets(..., collections="tag1")], 

256 [self.comparableRef(ref) for ref in registry2.queryDatasets(..., collections="tag1")], 

257 ) 

258 # Check that calibration collection contents are the same. 

259 self.assertCountEqual( 

260 [ 

261 (self.comparableRef(assoc.ref), assoc.timespan) 

262 for assoc in registry1.queryDatasetAssociations("bias", collections="calibration1") 

263 ], 

264 [ 

265 (self.comparableRef(assoc.ref), assoc.timespan) 

266 for assoc in registry2.queryDatasetAssociations("bias", collections="calibration1") 

267 ], 

268 ) 

269 

270 def testButlerGet(self): 

271 """Test that butler.get can work with different variants.""" 

272 # Import data to play with. 

273 butler = self.makeButler(writeable=True) 

274 butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml") 

275 butler.import_(filename=f"resource://lsst.daf.butler/tests/registry_data/{self.datasetsImportFile}") 

276 

277 # Find the DatasetRef for a flat 

278 coll = "imported_g" 

279 flat2g = butler.find_dataset( 

280 "flat", instrument="Cam1", full_name="Ab", physical_filter="Cam1-G", collections=coll 

281 ) 

282 

283 # Create a numpy integer to check that works fine 

284 detector_np = np.int64(2) if np else 2 

285 

286 # Try to get it using different variations of dataId + keyword 

287 # arguments 

288 # Note that instrument.class_name does not work 

289 variants = ( 

290 (None, {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}), 

291 (None, {"instrument": "Cam1", "detector": detector_np, "physical_filter": "Cam1-G"}), 

292 ({"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}, {}), 

293 ({"instrument": "Cam1", "detector": detector_np, "physical_filter": "Cam1-G"}, {}), 

294 ({"instrument": "Cam1", "detector": 2}, {"physical_filter": "Cam1-G"}), 

295 ({"detector.full_name": "Ab"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

296 ({"full_name": "Ab"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

297 (None, {"full_name": "Ab", "instrument": "Cam1", "physical_filter": "Cam1-G"}), 

298 (None, {"detector": "Ab", "instrument": "Cam1", "physical_filter": "Cam1-G"}), 

299 ({"name_in_raft": "b", "raft": "A"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

300 ({"name_in_raft": "b"}, {"raft": "A", "instrument": "Cam1", "physical_filter": "Cam1-G"}), 

301 (None, {"name_in_raft": "b", "raft": "A", "instrument": "Cam1", "physical_filter": "Cam1-G"}), 

302 ( 

303 {"detector.name_in_raft": "b", "detector.raft": "A"}, 

304 {"instrument": "Cam1", "physical_filter": "Cam1-G"}, 

305 ), 

306 ( 

307 { 

308 "detector.name_in_raft": "b", 

309 "detector.raft": "A", 

310 "instrument": "Cam1", 

311 "physical_filter": "Cam1-G", 

312 }, 

313 {}, 

314 ), 

315 # Duplicate (but valid) information. 

316 (None, {"instrument": "Cam1", "detector": 2, "raft": "A", "physical_filter": "Cam1-G"}), 

317 ({"detector": 2}, {"instrument": "Cam1", "raft": "A", "physical_filter": "Cam1-G"}), 

318 ({"raft": "A"}, {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}), 

319 ({"raft": "A"}, {"instrument": "Cam1", "detector": "Ab", "physical_filter": "Cam1-G"}), 

320 ) 

321 butler._metrics.reset() 

322 n_got = 0 

323 for dataId, kwds in variants: 

324 try: 

325 flat_id, _ = butler.get("flat", dataId=dataId, collections=coll, **kwds) 

326 n_got += 1 

327 except Exception as e: 

328 e.add_note(f"dataId={dataId}, kwds={kwds}") 

329 raise 

330 self.assertEqual(flat_id, flat2g.id, msg=f"DataId: {dataId}, kwds: {kwds}") 

331 self.assertEqual(butler._metrics.n_get, n_got) 

332 

333 # Check that bad combinations raise. 

334 variants = ( 

335 # Inconsistent detector information. 

336 (None, {"instrument": "Cam1", "detector": 2, "raft": "B", "physical_filter": "Cam1-G"}), 

337 ({"detector": 2}, {"instrument": "Cam1", "raft": "B", "physical_filter": "Cam1-G"}), 

338 ({"detector": 12}, {"instrument": "Cam1", "raft": "B", "physical_filter": "Cam1-G"}), 

339 ({"raft": "B"}, {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}), 

340 ({"raft": "B"}, {"instrument": "Cam1", "detector": "Ab", "physical_filter": "Cam1-G"}), 

341 # Under-specified. 

342 ({"raft": "B"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

343 # Spurious kwargs. 

344 (None, {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G", "x": "y"}), 

345 ({"x": "y"}, {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}), 

346 ) 

347 for dataId, kwds in variants: 

348 with self.assertRaises((ValueError, LookupError)): 

349 butler.get("flat", dataId=dataId, collections=coll, **kwds) 

350 

351 def testGetCalibration(self): 

352 """Test that `Butler.get` can be used to fetch from 

353 `~CollectionType.CALIBRATION` collections if the data ID includes 

354 extra dimensions with temporal information. 

355 """ 

356 # Import data to play with. 

357 butler = self.makeButler(writeable=True) 

358 butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml") 

359 butler.import_(filename=f"resource://lsst.daf.butler/tests/registry_data/{self.datasetsImportFile}") 

360 # Certify some biases into a CALIBRATION collection. 

361 registry = butler.registry 

362 registry.registerCollection("calibs", CollectionType.CALIBRATION) 

363 t1 = astropy.time.Time("2020-01-01T01:00:00", format="isot", scale="tai") 

364 t2 = astropy.time.Time("2020-01-01T02:00:00", format="isot", scale="tai") 

365 t3 = astropy.time.Time("2020-01-01T03:00:00", format="isot", scale="tai") 

366 bias1a = registry.findDataset("bias", instrument="Cam1", detector=1, collections="imported_g") 

367 bias2a = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g") 

368 bias3a = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g") 

369 bias2b = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r") 

370 bias3b = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r") 

371 registry.certify("calibs", [bias1a], Timespan(t1, t2)) 

372 registry.certify("calibs", [bias2a, bias3a], Timespan(t1, t2)) 

373 registry.certify("calibs", [bias2b], Timespan(t2, None)) 

374 registry.certify("calibs", [bias3b], Timespan(t2, t3)) 

375 # Insert some exposure dimension data. 

376 registry.insertDimensionData( 

377 "group", 

378 {"instrument": "Cam1", "group": "three"}, 

379 {"instrument": "Cam1", "group": "four"}, 

380 ) 

381 registry.insertDimensionData( 

382 "day_obs", 

383 {"instrument": "Cam1", "id": 20211114}, 

384 ) 

385 # Choose timespans for exposures within the above calibration ranges 

386 # but make sure they are not identical to the full range. 

387 exp_time = astropy.time.TimeDelta(15.0, format="sec", scale="tai") 

388 span_delta = t2 - t1 

389 exp3_begin = t1 + (span_delta / 2.0) 

390 exp3_end = exp3_begin + exp_time 

391 span_delta = t3 - t2 

392 exp4_begin = t2 + (span_delta / 2.0) 

393 exp4_end = exp4_begin + exp_time 

394 registry.insertDimensionData( 

395 "exposure", 

396 { 

397 "instrument": "Cam1", 

398 "id": 3, 

399 "obs_id": "three", 

400 "timespan": Timespan(exp3_begin, exp3_end), 

401 "physical_filter": "Cam1-G", 

402 "group": "three", 

403 "day_obs": 20211114, 

404 "seq_num": 55, 

405 }, 

406 { 

407 "instrument": "Cam1", 

408 "id": 4, 

409 "obs_id": "four", 

410 "timespan": Timespan(exp4_begin, exp4_end), 

411 "physical_filter": "Cam1-G", 

412 "group": "four", 

413 "day_obs": 20211114, 

414 "seq_num": 42, 

415 }, 

416 ) 

417 # Get some biases from raw-like data IDs. 

418 bias2a_id, _ = butler.get( 

419 "bias", {"instrument": "Cam1", "exposure": 3, "detector": 2}, collections="calibs" 

420 ) 

421 self.assertEqual(bias2a_id, bias2a.id) 

422 bias3b_id, _ = butler.get( 

423 "bias", {"instrument": "Cam1", "exposure": 4, "detector": 3}, collections="calibs" 

424 ) 

425 self.assertEqual(bias3b_id, bias3b.id) 

426 

427 # Use explicit timespan and no exposure record. 

428 bias3b_id, _ = butler.get( 

429 "bias", 

430 {"instrument": "Cam1", "detector": 3}, 

431 collections="calibs", 

432 timespan=Timespan(exp4_begin, exp4_end), 

433 ) 

434 self.assertEqual(bias3b_id, bias3b.id) 

435 

436 # No timespan at all. 

437 # Only one matching dataset in calibs collection so this works with 

438 # a defaulted timespan. 

439 bias1a_id, _ = butler.get("bias", {"instrument": "Cam1", "detector": 1}, collections="calibs") 

440 self.assertEqual(bias1a_id, bias1a.id) 

441 

442 # Multiple datasets match in calibs collection with infinite timespan 

443 # so this fails. 

444 with self.assertRaises(LookupError): 

445 bias3b_id, _ = butler.get("bias", {"instrument": "Cam1", "detector": 3}, collections="calibs") 

446 

447 # Get using the kwarg form 

448 bias3b_id, _ = butler.get("bias", instrument="Cam1", exposure=4, detector=3, collections="calibs") 

449 self.assertEqual(bias3b_id, bias3b.id) 

450 

451 # Do it again but using the record information 

452 bias2a_id, _ = butler.get( 

453 "bias", 

454 {"instrument": "Cam1", "exposure.obs_id": "three", "detector.full_name": "Ab"}, 

455 collections="calibs", 

456 ) 

457 self.assertEqual(bias2a_id, bias2a.id) 

458 bias3b_id, _ = butler.get( 

459 "bias", 

460 {"exposure.obs_id": "four", "detector.full_name": "Ba"}, 

461 collections="calibs", 

462 instrument="Cam1", 

463 ) 

464 self.assertEqual(bias3b_id, bias3b.id) 

465 

466 # And again but this time using the alternate value rather than 

467 # the primary. 

468 bias3b_id, _ = butler.get( 

469 "bias", {"exposure": "four", "detector": "Ba"}, collections="calibs", instrument="Cam1" 

470 ) 

471 self.assertEqual(bias3b_id, bias3b.id) 

472 

473 # And again but this time using the alternate value rather than 

474 # the primary and do it in the keyword arguments. 

475 bias3b_id, _ = butler.get( 

476 "bias", exposure="four", detector="Ba", collections="calibs", instrument="Cam1" 

477 ) 

478 self.assertEqual(bias3b_id, bias3b.id) 

479 

480 # Now with implied record columns 

481 bias3b_id, _ = butler.get( 

482 "bias", 

483 day_obs=20211114, 

484 seq_num=42, 

485 raft="B", 

486 name_in_raft="a", 

487 collections="calibs", 

488 instrument="Cam1", 

489 ) 

490 self.assertEqual(bias3b_id, bias3b.id) 

491 

492 # Allow a fully-specified dataId and unnecessary extra information 

493 # that comes from the record. 

494 bias3b_id, _ = butler.get( 

495 "bias", 

496 dataId=dict( 

497 exposure=4, 

498 day_obs=20211114, 

499 seq_num=42, 

500 detector=3, 

501 instrument="Cam1", 

502 ), 

503 collections="calibs", 

504 ) 

505 self.assertEqual(bias3b_id, bias3b.id) 

506 

507 # Query for a calibration in a RUN and CALIBRATION collection to 

508 # ensure we do not get duplicate results. 

509 results = butler.query_datasets("bias", collections=["calibs", "imported_g"], find_first=False) 

510 self.assertEqual(len(set(results)), len(results)) 

511 

512 # Extra but inconsistent record values are a problem. 

513 with self.assertRaises(ValueError): 

514 bias3b_id, _ = butler.get( 

515 "bias", 

516 exposure=3, 

517 day_obs=20211114, 

518 seq_num=42, 

519 detector=3, 

520 collections="calibs", 

521 instrument="Cam1", 

522 ) 

523 

524 # Ensure that spurious kwargs cause an exception. 

525 with self.assertRaises(ValueError): 

526 butler.get( 

527 "bias", 

528 {"exposure.obs_id": "four", "immediate": True, "detector.full_name": "Ba"}, 

529 collections="calibs", 

530 instrument="Cam1", 

531 ) 

532 

533 with self.assertRaises(ValueError): 

534 butler.get( 

535 "bias", 

536 day_obs=20211114, 

537 seq_num=42, 

538 raft="B", 

539 name_in_raft="a", 

540 collections="calibs", 

541 instrument="Cam1", 

542 immediate=True, 

543 ) 

544 

545 def testRegistryDefaults(self): 

546 """Test that we can default the collections and some data ID keys when 

547 constructing a butler. 

548 

549 Many tests that use default run already exist in ``test_butler.py``, so 

550 that isn't tested here. And while most of this functionality is 

551 implemented in `Registry`, we test it here instead of 

552 ``daf/butler/tests/registry.py`` because it shouldn't depend on the 

553 database backend at all. 

554 """ 

555 butler = self.makeButler(writeable=True) 

556 butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml") 

557 butler.import_(filename=f"resource://lsst.daf.butler/tests/registry_data/{self.datasetsImportFile}") 

558 # Need to actually set defaults later, not at construction, because 

559 # we need to import the instrument before we can use it as a default. 

560 # Don't set a default instrument value for data IDs, because 'Cam1' 

561 # should be inferred by virtue of that being the only value in the 

562 # input collections. 

563 butler.registry.defaults = RegistryDefaults(collections=["imported_g"]) 

564 # Use findDataset without collections or instrument. 

565 ref = butler.find_dataset("flat", detector=2, physical_filter="Cam1-G") 

566 # Do the same with Butler.get; this should ultimately invoke a lot of 

567 # the same code, so it's a bit circular, but mostly we're checking that 

568 # it works at all. 

569 dataset_id, _ = butler.get("flat", detector=2, physical_filter="Cam1-G") 

570 self.assertEqual(ref.id, dataset_id) 

571 # Query for datasets. Test defaulting the data ID in both kwargs and 

572 # in the WHERE expression. 

573 queried_refs_1 = set(butler.registry.queryDatasets("flat", detector=2, physical_filter="Cam1-G")) 

574 self.assertEqual({ref}, queried_refs_1) 

575 queried_refs_2 = set( 

576 butler.registry.queryDatasets("flat", where="detector=2 AND physical_filter='Cam1-G'") 

577 ) 

578 self.assertEqual({ref}, queried_refs_2) 

579 # Query for data IDs with a dataset constraint. 

580 queried_data_ids = set( 

581 butler.registry.queryDataIds( 

582 {"instrument", "detector", "physical_filter"}, 

583 datasets={"flat"}, 

584 detector=2, 

585 physical_filter="Cam1-G", 

586 ) 

587 ) 

588 self.assertEqual({ref.dataId}, queried_data_ids) 

589 # Add another instrument to the repo, and a dataset that uses it to 

590 # the `imported_g` collection. 

591 butler.registry.insertDimensionData("instrument", {"name": "Cam2"}) 

592 camera = DatasetType( 

593 "camera", 

594 dimensions=butler.dimensions["instrument"].minimal_group, 

595 storageClass="Camera", 

596 ) 

597 butler.registry.registerDatasetType(camera) 

598 butler.registry.insertDatasets(camera, [{"instrument": "Cam2"}], run="imported_g") 

599 # Initialize a new butler with `imported_g` as its default run. 

600 # This should not have a default instrument, because there are two. 

601 # Pass run instead of collections; this should set both. 

602 butler2 = Butler.from_config(butler=butler, run="imported_g") 

603 self.assertEqual(list(butler2.registry.defaults.collections), ["imported_g"]) 

604 self.assertEqual(butler2.registry.defaults.run, "imported_g") 

605 self.assertFalse(butler2.registry.defaults.dataId) 

606 # Initialize a new butler with an instrument default explicitly given. 

607 # Set collections instead of run, which should then be None. 

608 butler3 = Butler.from_config(butler=butler, collections=["imported_g"], instrument="Cam2") 

609 self.assertEqual(list(butler3.registry.defaults.collections), ["imported_g"]) 

610 self.assertIsNone(butler3.registry.defaults.run, None) 

611 self.assertEqual(butler3.registry.defaults.dataId.required, {"instrument": "Cam2"}) 

612 

613 # Check that repr() does not fail. 

614 defaults = RegistryDefaults(collections=["imported_g"], run="test") 

615 r = repr(defaults) 

616 self.assertIn("collections=('imported_g',)", r) 

617 self.assertIn("run='test'", r) 

618 

619 defaults = RegistryDefaults(run="test", instrument="DummyCam", skypix="pix") 

620 r = repr(defaults) 

621 self.assertIn("skypix='pix'", r) 

622 self.assertIn("instrument='DummyCam'", r) 

623 

624 def testJson(self): 

625 """Test JSON serialization mediated by registry.""" 

626 butler = self.makeButler(writeable=True) 

627 butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml") 

628 butler.import_(filename=f"resource://lsst.daf.butler/tests/registry_data/{self.datasetsImportFile}") 

629 # Need to actually set defaults later, not at construction, because 

630 # we need to import the instrument before we can use it as a default. 

631 # Don't set a default instrument value for data IDs, because 'Cam1' 

632 # should be inferred by virtue of that being the only value in the 

633 # input collections. 

634 butler.registry.defaults = RegistryDefaults(collections=["imported_g"]) 

635 # Use findDataset without collections or instrument. 

636 ref = butler.find_dataset("flat", detector=2, physical_filter="Cam1-G") 

637 

638 # Transform the ref and dataset type to and from JSON 

639 # and check that it can be reconstructed properly 

640 

641 # Do it with the ref and a component ref in minimal and standard form 

642 compRef = ref.makeComponentRef("wcs") 

643 

644 for test_item in (ref, ref.datasetType, compRef, compRef.datasetType): 

645 for minimal in (False, True): 

646 json_str = test_item.to_json(minimal=minimal) 

647 from_json = type(test_item).from_json(json_str, registry=butler.registry) 

648 self.assertEqual(from_json, test_item, msg=f"From JSON '{json_str}' using registry") 

649 

650 # for minimal=False case also do a test without registry 

651 if not minimal: 

652 from_json = type(test_item).from_json(json_str, universe=butler.dimensions) 

653 self.assertEqual(from_json, test_item, msg=f"From JSON '{json_str}' using universe") 

654 

655 def test_populated_by(self): 

656 """Test that dimension records can find other records.""" 

657 butler = self.makeButler(writeable=True) 

658 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "hsc-rc2-subset.yaml")) 

659 

660 elements = frozenset(element for element in butler.dimensions.elements if element.has_own_table) 

661 

662 # Get a visit-based dataId. 

663 data_ids = set(butler.registry.queryDataIds("visit", visit=1232, instrument="HSC")) 

664 

665 # Request all the records related to it. 

666 records = butler._extract_all_dimension_records_from_data_ids(butler, data_ids, elements) 

667 

668 self.assertIn(butler.dimensions["visit_detector_region"], records, f"Keys: {records.keys()}") 

669 self.assertIn(butler.dimensions["visit_system_membership"], records) 

670 self.assertIn(butler.dimensions["visit_system"], records) 

671 

672 def testJsonDimensionRecordsAndHtmlRepresentation(self): 

673 # Dimension Records 

674 butler = self.makeButler(writeable=True) 

675 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "hsc-rc2-subset.yaml")) 

676 

677 for dimension in ("detector", "visit", "exposure", "day_obs", "group"): 

678 records = butler.registry.queryDimensionRecords(dimension, instrument="HSC") 

679 for r in records: 

680 for minimal in (True, False): 

681 json_str = r.to_json(minimal=minimal) 

682 r_json = type(r).from_json(json_str, registry=butler.registry) 

683 self.assertEqual(r_json, r) 

684 # check with direct method 

685 simple = r.to_simple() 

686 fromDirect = type(simple).direct(**json.loads(json_str)) 

687 self.assertEqual(simple, fromDirect) 

688 # Also check equality of each of the components as dicts 

689 self.assertEqual(r_json.toDict(), r.toDict()) 

690 

691 # check the html representation of records 

692 r_html = r._repr_html_() 

693 self.assertTrue(isinstance(r_html, str)) 

694 self.assertIn(dimension, r_html) 

695 

696 def test_dimension_records_import(self): 

697 # Dimension Records 

698 butler = self.makeButler(writeable=True) 

699 with self.assertWarns(UserWarning) as cm: 

700 butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/hsc-rc2-subset-v0.yaml") 

701 self.assertIn("Constructing day_obs records with no timespans", str(cm.warning)) 

702 

703 # Count records and assume this means it worked. 

704 dimensions = ( 

705 ("day_obs", 3), 

706 ("group", 1), 

707 ("exposure", 1), 

708 ("visit", 2), 

709 ("detector", 3), 

710 ("visit_system_membership", 2), 

711 ) 

712 for dimension, count in dimensions: 

713 records = list(butler.registry.queryDimensionRecords(dimension, instrument="HSC")) 

714 self.assertEqual(len(records), count, dimension) 

715 

716 def testWildcardQueries(self): 

717 """Test that different collection type queries work.""" 

718 # Import data to play with. 

719 butler = self.makeButler(writeable=True) 

720 butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml") 

721 

722 # Create some collections 

723 created = {"collection", "u/user/test", "coll3"} 

724 for collection in created: 

725 butler.registry.registerCollection(collection, type=CollectionType.RUN) 

726 

727 collections = butler.registry.queryCollections() 

728 self.assertEqual(set(collections), created) 

729 

730 expressions = [ 

731 ("collection", {"collection"}), 

732 (..., created), 

733 ("*", created), 

734 (("collection", "*"), created), 

735 ("u/*", {"u/user/test"}), 

736 ("*oll*", {"collection", "coll3"}), 

737 ("*[0-9]", {"coll3"}), 

738 ] 

739 for expression, expected in expressions: 

740 result = butler.registry.queryCollections(expression) 

741 self.assertEqual(set(result), expected) 

742 

743 if self.supportsCollectionRegex: 

744 expressions = [ 

745 (re.compile("u.*"), {"u/user/test"}), 

746 (re.compile(".*oll.*"), {"collection", "coll3"}), 

747 ((re.compile(r".*\d$"), "u/user/test"), {"coll3", "u/user/test"}), 

748 ] 

749 for expression, expected in expressions: 

750 with self.assertWarns(FutureWarning): 

751 result = butler.registry.queryCollections(expression) 

752 self.assertEqual(set(result), expected) 

753 

754 def test_skypix_templates(self): 

755 """Test that skypix templates can work.""" 

756 # Dimension Records 

757 butler = self.makeButler(writeable=True) 

758 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "hsc-rc2-subset.yaml")) 

759 

760 sc = StorageClass("null") 

761 dataset_type = DatasetType("warp", ("visit", "htm7"), sc, universe=butler.dimensions) 

762 dataId = butler.registry.expandDataId( 

763 DataCoordinate.standardize( 

764 dict(visit=27136, htm7=12345, instrument="HSC"), universe=butler.dimensions 

765 ) 

766 ) 

767 ref = DatasetRef(dataset_type, dataId, run="test") 

768 self.assertTrue(ref.dataId.hasRecords()) 

769 

770 tmplstr = ( 

771 "{run}/{datasetType}/{visit.name|exposure.obs_id|xyz}_" 

772 "{instrument}_{skypix}_{htm7}_{skypix.id}_{htm7.id}" 

773 ) 

774 file_template = FileTemplate(tmplstr) 

775 file_template.validateTemplate(ref) 

776 path = file_template.format(ref) 

777 self.assertEqual(path, "test/warp/HSCA02713600_HSC_12345_12345_12345_12345") 

778 

779 def test_clone(self): 

780 # This just tests that the default-overriding logic works as expected. 

781 # The actual internals are tested in test_butler.py, in 

782 # ClonedSqliteButlerTestCase and 

783 # ClonedPostgresPosixDatastoreButlerTestCase. 

784 

785 butler = self.makeButler(writeable=True) 

786 butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml") 

787 butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/datasets.yaml") 

788 butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/spatial.yaml") 

789 

790 # Original butler was created with the default arguments: 

791 # collections = None 

792 # run = None 

793 # inferDefaults = True 

794 # no explicit default data ID 

795 

796 # Collections can be overridden, and default data ID will be inferred 

797 # from it. 

798 clone1 = butler.clone(collections="imported_g") 

799 self.assertEqual(clone1.registry.defaults.dataId, {"instrument": "Cam1"}) 

800 self.assertCountEqual(clone1.registry.defaults.collections, ["imported_g"]) 

801 self.assertIsNone(clone1.run) 

802 

803 # Disabling inferDefaults stops default data ID from being inferred 

804 # from collections. 

805 clone2 = clone1.clone(inferDefaults=False) 

806 self.assertEqual(clone2.registry.defaults.dataId, {}) 

807 self.assertCountEqual(clone2.registry.defaults.collections, ["imported_g"]) 

808 self.assertIsNone(clone2.run) 

809 

810 # Setting a new run doesn't override explicitly-set collections. 

811 clone3 = clone2.clone(run="imported_r") 

812 self.assertEqual(clone3.registry.defaults.dataId, {}) 

813 self.assertCountEqual(clone3.registry.defaults.collections, ["imported_g"]) 

814 self.assertEqual(clone3.run, "imported_r") 

815 

816 # Following the behavior of the Butler() constructor, run will populate 

817 # collections if collections was None. Default data ID is inferred 

818 # from the run collection. 

819 clone4 = butler.clone(run="imported_r") 

820 self.assertEqual(clone4.registry.defaults.dataId, {"instrument": "Cam1"}) 

821 self.assertCountEqual(clone4.registry.defaults.collections, ["imported_r"]) 

822 self.assertEqual(clone4.run, "imported_r") 

823 

824 # Explicitly set data ID is combined with inferred defaults from 

825 # collections. 

826 clone5 = clone4.clone(dataId={"skymap": "SkyMap1"}) 

827 self.assertEqual(clone5.registry.defaults.dataId, {"instrument": "Cam1", "skymap": "SkyMap1"}) 

828 self.assertCountEqual(clone5.registry.defaults.collections, ["imported_r"]) 

829 self.assertEqual(clone5.run, "imported_r") 

830 

831 # Disabling inferred defaults preserves explicitly set data ID 

832 clone6 = clone5.clone(inferDefaults=False) 

833 self.assertEqual(clone6.registry.defaults.dataId, {"skymap": "SkyMap1"}) 

834 self.assertCountEqual(clone5.registry.defaults.collections, ["imported_r"]) 

835 self.assertEqual(clone5.run, "imported_r") 

836 

837 def test_calibration_dataset_type_registration(self) -> None: 

838 # Register two dataset types that should share the same tags table, 

839 # but only one is a calibration and hence needs a calibs table. 

840 butler1 = self.makeButler(writeable=True) 

841 a = DatasetType("a", ["instrument"], universe=butler1.dimensions, storageClass="StructuredDataDict") 

842 b = DatasetType( 

843 "b", 

844 ["instrument"], 

845 universe=butler1.dimensions, 

846 storageClass="StructuredDataDict", 

847 isCalibration=True, 

848 ) 

849 butler1.registry.registerDatasetType(a) 

850 butler1.registry.registerDatasetType(b) 

851 self.assertEqual(butler1.get_dataset_type("a"), a) 

852 self.assertEqual(butler1.get_dataset_type("b"), b) 

853 butler1.registry.refresh() 

854 self.assertEqual(butler1.get_dataset_type("a"), a) 

855 self.assertEqual(butler1.get_dataset_type("b"), b) 

856 # Register them in the opposite order in a new repo. 

857 butler2 = self.makeButler(writeable=True) 

858 # Dataset types have to use correct universe and with RemoteButler 

859 # each butler instance has its own universe instance. 

860 a = DatasetType("a", ["instrument"], universe=butler2.dimensions, storageClass="StructuredDataDict") 

861 b = DatasetType( 

862 "b", 

863 ["instrument"], 

864 universe=butler2.dimensions, 

865 storageClass="StructuredDataDict", 

866 isCalibration=True, 

867 ) 

868 butler2.registry.registerDatasetType(b) 

869 butler2.registry.registerDatasetType(a) 

870 self.assertEqual(butler2.get_dataset_type("a"), a) 

871 self.assertEqual(butler2.get_dataset_type("b"), b) 

872 butler2.registry.refresh() 

873 self.assertEqual(butler2.get_dataset_type("a"), a) 

874 self.assertEqual(butler2.get_dataset_type("b"), b) 

875 

876 

877class DirectSimpleButlerTestCase(SimpleButlerTests, unittest.TestCase): 

878 """Run tests against DirectButler implementation.""" 

879 

880 datasetsManager = ( 

881 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID" 

882 ) 

883 

884 collectionsManager = "lsst.daf.butler.registry.collections.synthIntKey.SynthIntKeyCollectionManager" 

885 

886 def setUp(self): 

887 self.root = makeTestTempDir(TESTDIR) 

888 

889 def tearDown(self): 

890 removeTestTempDir(self.root) 

891 

892 def makeButler(self, writeable: bool = False) -> Butler: 

893 config = ButlerConfig() 

894 

895 # make separate temporary directory for registry of this instance 

896 tmpdir = tempfile.mkdtemp(dir=self.root) 

897 config["registry", "db"] = f"sqlite:///{tmpdir}/gen3.sqlite3" 

898 config["registry", "managers", "datasets"] = self.datasetsManager 

899 config["registry", "managers", "collections"] = self.collectionsManager 

900 config["root"] = self.root 

901 

902 # have to make a registry first 

903 registryConfig = RegistryConfig(config.get("registry")) 

904 _RegistryFactory(registryConfig).create_from_config() 

905 

906 # Write the YAML file so that some tests can recreate butler from it. 

907 config.dumpToUri(os.path.join(self.root, "butler.yaml")) 

908 butler = Butler.from_config(config, writeable=writeable) 

909 DatastoreMock.apply(butler) 

910 return butler 

911 

912 def test_dataset_uris(self): 

913 """Test that dataset URIs can be parsed and retrieved.""" 

914 butler = self.makeButler(writeable=True) 

915 butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml") 

916 butler.import_(filename=f"resource://lsst.daf.butler/tests/registry_data/{self.datasetsImportFile}") 

917 

918 butler.registry.defaults = RegistryDefaults(collections=["imported_g"]) 

919 ref = butler.find_dataset("flat", detector=2, physical_filter="Cam1-G") 

920 self.assertIsInstance(ref, DatasetRef) 

921 

922 # Get the butler root for the URI. It does have to be encoded 

923 # in case there are special characters in the path. 

924 config_dir = urllib.parse.quote(butler._config["root"]) 

925 

926 # Read it via a repo label and a path. 

927 with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml") as index_file: 

928 label = "test_repo" 

929 index_file.write(f"{label}: {config_dir}\n") 

930 index_file.flush() 

931 with mock_env({"DAF_BUTLER_REPOSITORY_INDEX": index_file.name}): 

932 butler_factory = LabeledButlerFactory() 

933 factory = butler_factory.bind(access_token=None) 

934 

935 for dataset_uri in ( 

936 f"ivo://org.rubinobs/usdac/test?repo={config_dir}&id={ref.id}", 

937 f"ivo://org.rubinobs/ukdac/lsst-dr1?repo={config_dir}%2Fbutler.yaml&id={ref.id}", 

938 f"butler://{label}/{ref.id}", 

939 f"ivo://org.rubinobs/usdac/lsst-dp1?repo={label}&id={ref.id}", 

940 ): 

941 result = Butler.get_dataset_from_uri(dataset_uri) 

942 self.assertEqual(result.dataset, ref) 

943 # The returned butler needs to have the datastore mocked. 

944 DatastoreMock.apply(result.butler) 

945 dataset_id, _ = result.butler.get(result.dataset) 

946 self.assertEqual(dataset_id, ref.id) 

947 

948 factory_result = Butler.get_dataset_from_uri(dataset_uri, factory=factory) 

949 self.assertEqual(factory_result.dataset, ref) 

950 # The returned butler needs to have the datastore mocked. 

951 DatastoreMock.apply(factory_result.butler) 

952 dataset_id, _ = factory_result.butler.get(factory_result.dataset) 

953 self.assertEqual(dataset_id, ref.id) 

954 

955 # Non existent dataset. 

956 missing_id = str(ref.id).replace("2", "3") 

957 result = Butler.get_dataset_from_uri(f"butler://{label}/{missing_id}") 

958 self.assertIsNone(result.dataset) 

959 

960 # Test some failure modes. 

961 for dataset_uri in ( 

962 "butler://label/1234", # Bad UUID. 

963 "butler://1234", # No UUID. 

964 "butler:///1234", # No label. 

965 "ivo://rubin/1234", # No query part and bad UUID and no label. 

966 "ivo://rubin/datasets/dr1/82d79caa-0823-4300-9874-67b737367ee0", # No query part. 

967 "ivo://org.rubinobs/datasets?repo=dr1&id=1234", # Bad UUID. 

968 "ivo://org.rubinobs/butler?release=dr1&id=82d79caa-0823-4300-9874-67b737367ee0", # No repo key. 

969 "ivo://org.rubinobs/butler?repo=dr1&repo=dr2&id=82d79caa-0823-4300-9874-67b737367ee0", # 2 vals. 

970 "ivo://org.rubinobs/something?repo=%20&id=82d79caa-0823-4300-9874-67b737367ee0", # no repo. 

971 "https://something.edu/1234", # Wrong scheme. 

972 ): 

973 with self.assertRaises(ValueError): 

974 Butler.parse_dataset_uri(dataset_uri) 

975 

976 

977class NameKeyCollectionManagerDirectSimpleButlerTestCase(DirectSimpleButlerTestCase, unittest.TestCase): 

978 """Run tests against DirectButler implementation using the 

979 NameKeyCollectionsManager. 

980 """ 

981 

982 collectionsManager = "lsst.daf.butler.registry.collections.nameKey.NameKeyCollectionManager" 

983 

984 

985@unittest.skipIf(create_test_server is None, "Server dependencies not installed.") 

986class RemoteSimpleButlerTestCase(SimpleButlerTests, unittest.TestCase): 

987 """Run tests against Butler client/server.""" 

988 

989 supportsCollectionRegex = False 

990 

991 def makeButler(self, writeable: bool = False) -> Butler: 

992 server_instance = self.enterContext(create_test_server(TESTDIR)) 

993 butler = server_instance.hybrid_butler 

994 DatastoreMock.apply(butler) 

995 # Because RemoteButler doesn't have a Datastore object, we have to 

996 # duplicate some of the functionality from DatastoreMock separately. 

997 butler._remote_butler._get_dataset_as_python_object = _mock_get_dataset_as_python_object 

998 return butler 

999 

1000 

1001def _mock_get_dataset_as_python_object( 

1002 ref: DatasetRef, 

1003 model: Any, 

1004 parameters: dict[str, Any] | None, 

1005) -> Any: 

1006 """Mimic the functionality of DatastoreMock's get() mock.""" 

1007 return (ref.id, parameters) 

1008 

1009 

1010if __name__ == "__main__": 

1011 unittest.main()