Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

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 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 <http://www.gnu.org/licenses/>. 

21from __future__ import annotations 

22 

23__all__ = ["RegistryTests"] 

24 

25from abc import ABC, abstractmethod 

26import itertools 

27import os 

28import re 

29import unittest 

30 

31import astropy.time 

32import sqlalchemy 

33from typing import Optional, Type, Union 

34 

35try: 

36 import numpy as np 

37except ImportError: 

38 np = None 

39 

40from ...core import ( 

41 DataCoordinate, 

42 DataCoordinateSequence, 

43 DataCoordinateSet, 

44 DatasetAssociation, 

45 DatasetRef, 

46 DatasetType, 

47 DimensionGraph, 

48 NamedValueSet, 

49 StorageClass, 

50 ddl, 

51 Timespan, 

52) 

53from .._registry import ( 

54 CollectionType, 

55 ConflictingDefinitionError, 

56 InconsistentDataIdError, 

57 Registry, 

58 RegistryConfig, 

59) 

60from ..wildcards import DatasetTypeRestriction 

61from ..interfaces import MissingCollectionError, ButlerAttributeExistsError 

62 

63 

64class RegistryTests(ABC): 

65 """Generic tests for the `Registry` class that can be subclassed to 

66 generate tests for different configurations. 

67 """ 

68 

69 collectionsManager: Optional[str] = None 

70 """Name of the collections manager class, if subclass provides value for 

71 this member then it overrides name specified in default configuration 

72 (`str`). 

73 """ 

74 

75 @classmethod 

76 @abstractmethod 

77 def getDataDir(cls) -> str: 

78 """Return the root directory containing test data YAML files. 

79 """ 

80 raise NotImplementedError() 

81 

82 def makeRegistryConfig(self) -> RegistryConfig: 

83 """Create RegistryConfig used to create a registry. 

84 

85 This method should be called by a subclass from `makeRegistry`. 

86 Returned instance will be pre-configured based on the values of class 

87 members, and default-configured for all other parametrs. Subclasses 

88 that need default configuration should just instantiate 

89 `RegistryConfig` directly. 

90 """ 

91 config = RegistryConfig() 

92 if self.collectionsManager: 

93 config["managers"]["collections"] = self.collectionsManager 

94 return config 

95 

96 @abstractmethod 

97 def makeRegistry(self) -> Registry: 

98 """Return the Registry instance to be tested. 

99 """ 

100 raise NotImplementedError() 

101 

102 def loadData(self, registry: Registry, filename: str): 

103 """Load registry test data from ``getDataDir/<filename>``, 

104 which should be a YAML import/export file. 

105 """ 

106 from ...transfers import YamlRepoImportBackend 

107 with open(os.path.join(self.getDataDir(), filename), 'r') as stream: 

108 backend = YamlRepoImportBackend(stream, registry) 

109 backend.register() 

110 backend.load(datastore=None) 

111 

112 def assertRowCount(self, registry: Registry, table: str, count: int): 

113 """Check the number of rows in table. 

114 """ 

115 # TODO: all tests that rely on this method should be rewritten, as it 

116 # needs to depend on Registry implementation details to have any chance 

117 # of working. 

118 sql = sqlalchemy.sql.select( 

119 [sqlalchemy.sql.func.count()] 

120 ).select_from( 

121 getattr(registry._tables, table) 

122 ) 

123 self.assertEqual(registry._db.query(sql).scalar(), count) 

124 

125 def testOpaque(self): 

126 """Tests for `Registry.registerOpaqueTable`, 

127 `Registry.insertOpaqueData`, `Registry.fetchOpaqueData`, and 

128 `Registry.deleteOpaqueData`. 

129 """ 

130 registry = self.makeRegistry() 

131 table = "opaque_table_for_testing" 

132 registry.registerOpaqueTable( 

133 table, 

134 spec=ddl.TableSpec( 

135 fields=[ 

136 ddl.FieldSpec("id", dtype=sqlalchemy.BigInteger, primaryKey=True), 

137 ddl.FieldSpec("name", dtype=sqlalchemy.String, length=16, nullable=False), 

138 ddl.FieldSpec("count", dtype=sqlalchemy.SmallInteger, nullable=True), 

139 ], 

140 ) 

141 ) 

142 rows = [ 

143 {"id": 1, "name": "one", "count": None}, 

144 {"id": 2, "name": "two", "count": 5}, 

145 {"id": 3, "name": "three", "count": 6}, 

146 ] 

147 registry.insertOpaqueData(table, *rows) 

148 self.assertCountEqual(rows, list(registry.fetchOpaqueData(table))) 

149 self.assertEqual(rows[0:1], list(registry.fetchOpaqueData(table, id=1))) 

150 self.assertEqual(rows[1:2], list(registry.fetchOpaqueData(table, name="two"))) 

151 self.assertEqual([], list(registry.fetchOpaqueData(table, id=1, name="two"))) 

152 registry.deleteOpaqueData(table, id=3) 

153 self.assertCountEqual(rows[:2], list(registry.fetchOpaqueData(table))) 

154 registry.deleteOpaqueData(table) 

155 self.assertEqual([], list(registry.fetchOpaqueData(table))) 

156 

157 def testDatasetType(self): 

158 """Tests for `Registry.registerDatasetType` and 

159 `Registry.getDatasetType`. 

160 """ 

161 registry = self.makeRegistry() 

162 # Check valid insert 

163 datasetTypeName = "test" 

164 storageClass = StorageClass("testDatasetType") 

165 registry.storageClasses.registerStorageClass(storageClass) 

166 dimensions = registry.dimensions.extract(("instrument", "visit")) 

167 differentDimensions = registry.dimensions.extract(("instrument", "patch")) 

168 inDatasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

169 # Inserting for the first time should return True 

170 self.assertTrue(registry.registerDatasetType(inDatasetType)) 

171 outDatasetType1 = registry.getDatasetType(datasetTypeName) 

172 self.assertEqual(outDatasetType1, inDatasetType) 

173 

174 # Re-inserting should work 

175 self.assertFalse(registry.registerDatasetType(inDatasetType)) 

176 # Except when they are not identical 

177 with self.assertRaises(ConflictingDefinitionError): 

178 nonIdenticalDatasetType = DatasetType(datasetTypeName, differentDimensions, storageClass) 

179 registry.registerDatasetType(nonIdenticalDatasetType) 

180 

181 # Template can be None 

182 datasetTypeName = "testNoneTemplate" 

183 storageClass = StorageClass("testDatasetType2") 

184 registry.storageClasses.registerStorageClass(storageClass) 

185 dimensions = registry.dimensions.extract(("instrument", "visit")) 

186 inDatasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

187 registry.registerDatasetType(inDatasetType) 

188 outDatasetType2 = registry.getDatasetType(datasetTypeName) 

189 self.assertEqual(outDatasetType2, inDatasetType) 

190 

191 allTypes = set(registry.queryDatasetTypes()) 

192 self.assertEqual(allTypes, {outDatasetType1, outDatasetType2}) 

193 

194 def testDimensions(self): 

195 """Tests for `Registry.insertDimensionData`, 

196 `Registry.syncDimensionData`, and `Registry.expandDataId`. 

197 """ 

198 registry = self.makeRegistry() 

199 dimensionName = "instrument" 

200 dimension = registry.dimensions[dimensionName] 

201 dimensionValue = {"name": "DummyCam", "visit_max": 10, "exposure_max": 10, "detector_max": 2, 

202 "class_name": "lsst.obs.base.Instrument"} 

203 registry.insertDimensionData(dimensionName, dimensionValue) 

204 # Inserting the same value twice should fail 

205 with self.assertRaises(sqlalchemy.exc.IntegrityError): 

206 registry.insertDimensionData(dimensionName, dimensionValue) 

207 # expandDataId should retrieve the record we just inserted 

208 self.assertEqual( 

209 registry.expandDataId( 

210 instrument="DummyCam", 

211 graph=dimension.graph 

212 ).records[dimensionName].toDict(), 

213 dimensionValue 

214 ) 

215 # expandDataId should raise if there is no record with the given ID. 

216 with self.assertRaises(LookupError): 

217 registry.expandDataId({"instrument": "Unknown"}, graph=dimension.graph) 

218 # band doesn't have a table; insert should fail. 

219 with self.assertRaises(TypeError): 

220 registry.insertDimensionData("band", {"band": "i"}) 

221 dimensionName2 = "physical_filter" 

222 dimension2 = registry.dimensions[dimensionName2] 

223 dimensionValue2 = {"name": "DummyCam_i", "band": "i"} 

224 # Missing required dependency ("instrument") should fail 

225 with self.assertRaises(KeyError): 

226 registry.insertDimensionData(dimensionName2, dimensionValue2) 

227 # Adding required dependency should fix the failure 

228 dimensionValue2["instrument"] = "DummyCam" 

229 registry.insertDimensionData(dimensionName2, dimensionValue2) 

230 # expandDataId should retrieve the record we just inserted. 

231 self.assertEqual( 

232 registry.expandDataId( 

233 instrument="DummyCam", physical_filter="DummyCam_i", 

234 graph=dimension2.graph 

235 ).records[dimensionName2].toDict(), 

236 dimensionValue2 

237 ) 

238 # Use syncDimensionData to insert a new record successfully. 

239 dimensionName3 = "detector" 

240 dimensionValue3 = {"instrument": "DummyCam", "id": 1, "full_name": "one", 

241 "name_in_raft": "zero", "purpose": "SCIENCE"} 

242 self.assertTrue(registry.syncDimensionData(dimensionName3, dimensionValue3)) 

243 # Sync that again. Note that one field ("raft") is NULL, and that 

244 # should be okay. 

245 self.assertFalse(registry.syncDimensionData(dimensionName3, dimensionValue3)) 

246 # Now try that sync with the same primary key but a different value. 

247 # This should fail. 

248 with self.assertRaises(ConflictingDefinitionError): 

249 registry.syncDimensionData( 

250 dimensionName3, 

251 {"instrument": "DummyCam", "id": 1, "full_name": "one", 

252 "name_in_raft": "four", "purpose": "SCIENCE"} 

253 ) 

254 

255 @unittest.skipIf(np is None, "numpy not available.") 

256 def testNumpyDataId(self): 

257 """Test that we can use a numpy int in a dataId.""" 

258 registry = self.makeRegistry() 

259 dimensionEntries = [ 

260 ("instrument", {"instrument": "DummyCam"}), 

261 ("physical_filter", {"instrument": "DummyCam", "name": "d-r", "band": "R"}), 

262 # Using an np.int64 here fails unless Records.fromDict is also 

263 # patched to look for numbers.Integral 

264 ("visit", {"instrument": "DummyCam", "id": 42, "name": "fortytwo", "physical_filter": "d-r"}), 

265 ] 

266 for args in dimensionEntries: 

267 registry.insertDimensionData(*args) 

268 

269 # Try a normal integer and something that looks like an int but 

270 # is not. 

271 for visit_id in (42, np.int64(42)): 

272 with self.subTest(visit_id=visit_id, id_type=type(visit_id).__name__): 

273 expanded = registry.expandDataId({"instrument": "DummyCam", "visit": visit_id}) 

274 self.assertEqual(expanded["visit"], int(visit_id)) 

275 self.assertIsInstance(expanded["visit"], int) 

276 

277 def testDataIdRelationships(self): 

278 """Test that `Registry.expandDataId` raises an exception when the given 

279 keys are inconsistent. 

280 """ 

281 registry = self.makeRegistry() 

282 self.loadData(registry, "base.yaml") 

283 # Insert a few more dimension records for the next test. 

284 registry.insertDimensionData( 

285 "exposure", 

286 {"instrument": "Cam1", "id": 1, "name": "one", "physical_filter": "Cam1-G"}, 

287 ) 

288 registry.insertDimensionData( 

289 "exposure", 

290 {"instrument": "Cam1", "id": 2, "name": "two", "physical_filter": "Cam1-G"}, 

291 ) 

292 registry.insertDimensionData( 

293 "visit_system", 

294 {"instrument": "Cam1", "id": 0, "name": "one-to-one"}, 

295 ) 

296 registry.insertDimensionData( 

297 "visit", 

298 {"instrument": "Cam1", "id": 1, "name": "one", "physical_filter": "Cam1-G", "visit_system": 0}, 

299 ) 

300 registry.insertDimensionData( 

301 "visit_definition", 

302 {"instrument": "Cam1", "visit": 1, "exposure": 1, "visit_system": 0}, 

303 ) 

304 with self.assertRaises(InconsistentDataIdError): 

305 registry.expandDataId( 

306 {"instrument": "Cam1", "visit": 1, "exposure": 2}, 

307 ) 

308 

309 def testDataset(self): 

310 """Basic tests for `Registry.insertDatasets`, `Registry.getDataset`, 

311 and `Registry.removeDatasets`. 

312 """ 

313 registry = self.makeRegistry() 

314 self.loadData(registry, "base.yaml") 

315 run = "test" 

316 registry.registerRun(run) 

317 datasetType = registry.getDatasetType("bias") 

318 dataId = {"instrument": "Cam1", "detector": 2} 

319 ref, = registry.insertDatasets(datasetType, dataIds=[dataId], run=run) 

320 outRef = registry.getDataset(ref.id) 

321 self.assertIsNotNone(ref.id) 

322 self.assertEqual(ref, outRef) 

323 with self.assertRaises(ConflictingDefinitionError): 

324 registry.insertDatasets(datasetType, dataIds=[dataId], run=run) 

325 registry.removeDatasets([ref]) 

326 self.assertIsNone(registry.findDataset(datasetType, dataId, collections=[run])) 

327 

328 def testFindDataset(self): 

329 """Tests for `Registry.findDataset`. 

330 """ 

331 registry = self.makeRegistry() 

332 self.loadData(registry, "base.yaml") 

333 run = "test" 

334 datasetType = registry.getDatasetType("bias") 

335 dataId = {"instrument": "Cam1", "detector": 4} 

336 registry.registerRun(run) 

337 inputRef, = registry.insertDatasets(datasetType, dataIds=[dataId], run=run) 

338 outputRef = registry.findDataset(datasetType, dataId, collections=[run]) 

339 self.assertEqual(outputRef, inputRef) 

340 # Check that retrieval with invalid dataId raises 

341 with self.assertRaises(LookupError): 

342 dataId = {"instrument": "Cam1"} # no detector 

343 registry.findDataset(datasetType, dataId, collections=run) 

344 # Check that different dataIds match to different datasets 

345 dataId1 = {"instrument": "Cam1", "detector": 1} 

346 inputRef1, = registry.insertDatasets(datasetType, dataIds=[dataId1], run=run) 

347 dataId2 = {"instrument": "Cam1", "detector": 2} 

348 inputRef2, = registry.insertDatasets(datasetType, dataIds=[dataId2], run=run) 

349 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=run), inputRef1) 

350 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=run), inputRef2) 

351 self.assertNotEqual(registry.findDataset(datasetType, dataId1, collections=run), inputRef2) 

352 self.assertNotEqual(registry.findDataset(datasetType, dataId2, collections=run), inputRef1) 

353 # Check that requesting a non-existing dataId returns None 

354 nonExistingDataId = {"instrument": "Cam1", "detector": 3} 

355 self.assertIsNone(registry.findDataset(datasetType, nonExistingDataId, collections=run)) 

356 

357 def testDatasetTypeComponentQueries(self): 

358 """Test component options when querying for dataset types. 

359 """ 

360 registry = self.makeRegistry() 

361 self.loadData(registry, "base.yaml") 

362 self.loadData(registry, "datasets.yaml") 

363 # Test querying for dataset types with different inputs. 

364 # First query for all dataset types; components should only be included 

365 # when components=True. 

366 self.assertEqual( 

367 {"bias", "flat"}, 

368 NamedValueSet(registry.queryDatasetTypes()).names 

369 ) 

370 self.assertEqual( 

371 {"bias", "flat"}, 

372 NamedValueSet(registry.queryDatasetTypes(components=False)).names 

373 ) 

374 self.assertLess( 

375 {"bias", "flat", "bias.wcs", "flat.photoCalib"}, 

376 NamedValueSet(registry.queryDatasetTypes(components=True)).names 

377 ) 

378 # Use a pattern that can match either parent or components. Again, 

379 # components are only returned if components=True. 

380 self.assertEqual( 

381 {"bias"}, 

382 NamedValueSet(registry.queryDatasetTypes(re.compile("^bias.*"))).names 

383 ) 

384 self.assertEqual( 

385 {"bias"}, 

386 NamedValueSet(registry.queryDatasetTypes(re.compile("^bias.*"), components=False)).names 

387 ) 

388 self.assertLess( 

389 {"bias", "bias.wcs"}, 

390 NamedValueSet(registry.queryDatasetTypes(re.compile("^bias.*"), components=True)).names 

391 ) 

392 # This pattern matches only a component. In this case we also return 

393 # that component dataset type if components=None. 

394 self.assertEqual( 

395 {"bias.wcs"}, 

396 NamedValueSet(registry.queryDatasetTypes(re.compile(r"^bias\.wcs"))).names 

397 ) 

398 self.assertEqual( 

399 set(), 

400 NamedValueSet(registry.queryDatasetTypes(re.compile(r"^bias\.wcs"), components=False)).names 

401 ) 

402 self.assertEqual( 

403 {"bias.wcs"}, 

404 NamedValueSet(registry.queryDatasetTypes(re.compile(r"^bias\.wcs"), components=True)).names 

405 ) 

406 

407 def testComponentLookups(self): 

408 """Test searching for component datasets via their parents. 

409 """ 

410 registry = self.makeRegistry() 

411 self.loadData(registry, "base.yaml") 

412 self.loadData(registry, "datasets.yaml") 

413 # Test getting the child dataset type (which does still exist in the 

414 # Registry), and check for consistency with 

415 # DatasetRef.makeComponentRef. 

416 collection = "imported_g" 

417 parentType = registry.getDatasetType("bias") 

418 childType = registry.getDatasetType("bias.wcs") 

419 parentRefResolved = registry.findDataset(parentType, collections=collection, 

420 instrument="Cam1", detector=1) 

421 self.assertIsInstance(parentRefResolved, DatasetRef) 

422 self.assertEqual(childType, parentRefResolved.makeComponentRef("wcs").datasetType) 

423 # Search for a single dataset with findDataset. 

424 childRef1 = registry.findDataset("bias.wcs", collections=collection, 

425 dataId=parentRefResolved.dataId) 

426 self.assertEqual(childRef1, parentRefResolved.makeComponentRef("wcs")) 

427 # Search for detector data IDs constrained by component dataset 

428 # existence with queryDataIds. 

429 dataIds = registry.queryDataIds( 

430 ["detector"], 

431 datasets=["bias.wcs"], 

432 collections=collection, 

433 ).toSet() 

434 self.assertEqual( 

435 dataIds, 

436 DataCoordinateSet( 

437 { 

438 DataCoordinate.standardize(instrument="Cam1", detector=d, graph=parentType.dimensions) 

439 for d in (1, 2, 3) 

440 }, 

441 parentType.dimensions, 

442 ) 

443 ) 

444 # Search for multiple datasets of a single type with queryDatasets. 

445 childRefs2 = set(registry.queryDatasets( 

446 "bias.wcs", 

447 collections=collection, 

448 )) 

449 self.assertEqual( 

450 {ref.unresolved() for ref in childRefs2}, 

451 {DatasetRef(childType, dataId) for dataId in dataIds} 

452 ) 

453 

454 def testCollections(self): 

455 """Tests for registry methods that manage collections. 

456 """ 

457 registry = self.makeRegistry() 

458 self.loadData(registry, "base.yaml") 

459 self.loadData(registry, "datasets.yaml") 

460 run1 = "imported_g" 

461 run2 = "imported_r" 

462 datasetType = "bias" 

463 # Find some datasets via their run's collection. 

464 dataId1 = {"instrument": "Cam1", "detector": 1} 

465 ref1 = registry.findDataset(datasetType, dataId1, collections=run1) 

466 self.assertIsNotNone(ref1) 

467 dataId2 = {"instrument": "Cam1", "detector": 2} 

468 ref2 = registry.findDataset(datasetType, dataId2, collections=run1) 

469 self.assertIsNotNone(ref2) 

470 # Associate those into a new collection,then look for them there. 

471 tag1 = "tag1" 

472 registry.registerCollection(tag1, type=CollectionType.TAGGED) 

473 registry.associate(tag1, [ref1, ref2]) 

474 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=tag1), ref1) 

475 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2) 

476 # Disassociate one and verify that we can't it there anymore... 

477 registry.disassociate(tag1, [ref1]) 

478 self.assertIsNone(registry.findDataset(datasetType, dataId1, collections=tag1)) 

479 # ...but we can still find ref2 in tag1, and ref1 in the run. 

480 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=run1), ref1) 

481 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2) 

482 collections = set(registry.queryCollections()) 

483 self.assertEqual(collections, {run1, run2, tag1}) 

484 # Associate both refs into tag1 again; ref2 is already there, but that 

485 # should be a harmless no-op. 

486 registry.associate(tag1, [ref1, ref2]) 

487 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=tag1), ref1) 

488 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2) 

489 # Get a different dataset (from a different run) that has the same 

490 # dataset type and data ID as ref2. 

491 ref2b = registry.findDataset(datasetType, dataId2, collections=run2) 

492 self.assertNotEqual(ref2, ref2b) 

493 # Attempting to associate that into tag1 should be an error. 

494 with self.assertRaises(ConflictingDefinitionError): 

495 registry.associate(tag1, [ref2b]) 

496 # That error shouldn't have messed up what we had before. 

497 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=tag1), ref1) 

498 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2) 

499 # Attempt to associate the conflicting dataset again, this time with 

500 # a dataset that isn't in the collection and won't cause a conflict. 

501 # Should also fail without modifying anything. 

502 dataId3 = {"instrument": "Cam1", "detector": 3} 

503 ref3 = registry.findDataset(datasetType, dataId3, collections=run1) 

504 with self.assertRaises(ConflictingDefinitionError): 

505 registry.associate(tag1, [ref3, ref2b]) 

506 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=tag1), ref1) 

507 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2) 

508 self.assertIsNone(registry.findDataset(datasetType, dataId3, collections=tag1)) 

509 # Register a chained collection that searches: 

510 # 1. 'tag1' 

511 # 2. 'run1', but only for the flat dataset 

512 # 3. 'run2' 

513 chain1 = "chain1" 

514 registry.registerCollection(chain1, type=CollectionType.CHAINED) 

515 self.assertIs(registry.getCollectionType(chain1), CollectionType.CHAINED) 

516 # Chained collection exists, but has no collections in it. 

517 self.assertFalse(registry.getCollectionChain(chain1)) 

518 # If we query for all collections, we should get the chained collection 

519 # only if we don't ask to flatten it (i.e. yield only its children). 

520 self.assertEqual(set(registry.queryCollections(flattenChains=False)), {tag1, run1, run2, chain1}) 

521 self.assertEqual(set(registry.queryCollections(flattenChains=True)), {tag1, run1, run2}) 

522 # Attempt to set its child collections to something circular; that 

523 # should fail. 

524 with self.assertRaises(ValueError): 

525 registry.setCollectionChain(chain1, [tag1, chain1]) 

526 # Add the child collections. 

527 registry.setCollectionChain(chain1, [tag1, (run1, "flat"), run2]) 

528 self.assertEqual( 

529 list(registry.getCollectionChain(chain1)), 

530 [(tag1, DatasetTypeRestriction.any), 

531 (run1, DatasetTypeRestriction.fromExpression("flat")), 

532 (run2, DatasetTypeRestriction.any)] 

533 ) 

534 # Searching for dataId1 or dataId2 in the chain should return ref1 and 

535 # ref2, because both are in tag1. 

536 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=chain1), ref1) 

537 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=chain1), ref2) 

538 # Now disassociate ref2 from tag1. The search (for bias) with 

539 # dataId2 in chain1 should then: 

540 # 1. not find it in tag1 

541 # 2. not look in tag2, because it's restricted to flat here 

542 # 3. find a different dataset in run2 

543 registry.disassociate(tag1, [ref2]) 

544 ref2b = registry.findDataset(datasetType, dataId2, collections=chain1) 

545 self.assertNotEqual(ref2b, ref2) 

546 self.assertEqual(ref2b, registry.findDataset(datasetType, dataId2, collections=run2)) 

547 # Look in the chain for a flat that is in run1; should get the 

548 # same ref as if we'd searched run1 directly. 

549 dataId3 = {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"} 

550 self.assertEqual(registry.findDataset("flat", dataId3, collections=chain1), 

551 registry.findDataset("flat", dataId3, collections=run1),) 

552 # Define a new chain so we can test recursive chains. 

553 chain2 = "chain2" 

554 registry.registerCollection(chain2, type=CollectionType.CHAINED) 

555 registry.setCollectionChain(chain2, [(run2, "bias"), chain1]) 

556 # Query for collections matching a regex. 

557 self.assertCountEqual( 

558 list(registry.queryCollections(re.compile("imported_."), flattenChains=False)), 

559 ["imported_r", "imported_g"] 

560 ) 

561 # Query for collections matching a regex or an explicit str. 

562 self.assertCountEqual( 

563 list(registry.queryCollections([re.compile("imported_."), "chain1"], flattenChains=False)), 

564 ["imported_r", "imported_g", "chain1"] 

565 ) 

566 # Search for bias with dataId1 should find it via tag1 in chain2, 

567 # recursing, because is not in run1. 

568 self.assertIsNone(registry.findDataset(datasetType, dataId1, collections=run2)) 

569 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=chain2), ref1) 

570 # Search for bias with dataId2 should find it in run2 (ref2b). 

571 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=chain2), ref2b) 

572 # Search for a flat that is in run2. That should not be found 

573 # at the front of chain2, because of the restriction to bias 

574 # on run2 there, but it should be found in at the end of chain1. 

575 dataId4 = {"instrument": "Cam1", "detector": 3, "physical_filter": "Cam1-R2"} 

576 ref4 = registry.findDataset("flat", dataId4, collections=run2) 

577 self.assertIsNotNone(ref4) 

578 self.assertEqual(ref4, registry.findDataset("flat", dataId4, collections=chain2)) 

579 # Deleting a collection that's part of a CHAINED collection is not 

580 # allowed, and is exception-safe. 

581 with self.assertRaises(Exception): 

582 registry.removeCollection(run2) 

583 self.assertEqual(registry.getCollectionType(run2), CollectionType.RUN) 

584 with self.assertRaises(Exception): 

585 registry.removeCollection(chain1) 

586 self.assertEqual(registry.getCollectionType(chain1), CollectionType.CHAINED) 

587 # Actually remove chain2, test that it's gone by asking for its type. 

588 registry.removeCollection(chain2) 

589 with self.assertRaises(MissingCollectionError): 

590 registry.getCollectionType(chain2) 

591 # Actually remove run2 and chain1, which should work now. 

592 registry.removeCollection(chain1) 

593 registry.removeCollection(run2) 

594 with self.assertRaises(MissingCollectionError): 

595 registry.getCollectionType(run2) 

596 with self.assertRaises(MissingCollectionError): 

597 registry.getCollectionType(chain1) 

598 # Remove tag1 as well, just to test that we can remove TAGGED 

599 # collections. 

600 registry.removeCollection(tag1) 

601 with self.assertRaises(MissingCollectionError): 

602 registry.getCollectionType(tag1) 

603 

604 def testBasicTransaction(self): 

605 """Test that all operations within a single transaction block are 

606 rolled back if an exception propagates out of the block. 

607 """ 

608 registry = self.makeRegistry() 

609 storageClass = StorageClass("testDatasetType") 

610 registry.storageClasses.registerStorageClass(storageClass) 

611 with registry.transaction(): 

612 registry.insertDimensionData("instrument", {"name": "Cam1", "class_name": "A"}) 

613 with self.assertRaises(ValueError): 

614 with registry.transaction(): 

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

616 raise ValueError("Oops, something went wrong") 

617 # Cam1 should exist 

618 self.assertEqual(registry.expandDataId(instrument="Cam1").records["instrument"].class_name, "A") 

619 # But Cam2 and Cam3 should both not exist 

620 with self.assertRaises(LookupError): 

621 registry.expandDataId(instrument="Cam2") 

622 with self.assertRaises(LookupError): 

623 registry.expandDataId(instrument="Cam3") 

624 

625 def testNestedTransaction(self): 

626 """Test that operations within a transaction block are not rolled back 

627 if an exception propagates out of an inner transaction block and is 

628 then caught. 

629 """ 

630 registry = self.makeRegistry() 

631 dimension = registry.dimensions["instrument"] 

632 dataId1 = {"instrument": "DummyCam"} 

633 dataId2 = {"instrument": "DummyCam2"} 

634 checkpointReached = False 

635 with registry.transaction(): 

636 # This should be added and (ultimately) committed. 

637 registry.insertDimensionData(dimension, dataId1) 

638 with self.assertRaises(sqlalchemy.exc.IntegrityError): 

639 with registry.transaction(savepoint=True): 

640 # This does not conflict, and should succeed (but not 

641 # be committed). 

642 registry.insertDimensionData(dimension, dataId2) 

643 checkpointReached = True 

644 # This should conflict and raise, triggerring a rollback 

645 # of the previous insertion within the same transaction 

646 # context, but not the original insertion in the outer 

647 # block. 

648 registry.insertDimensionData(dimension, dataId1) 

649 self.assertTrue(checkpointReached) 

650 self.assertIsNotNone(registry.expandDataId(dataId1, graph=dimension.graph)) 

651 with self.assertRaises(LookupError): 

652 registry.expandDataId(dataId2, graph=dimension.graph) 

653 

654 def testInstrumentDimensions(self): 

655 """Test queries involving only instrument dimensions, with no joins to 

656 skymap.""" 

657 registry = self.makeRegistry() 

658 

659 # need a bunch of dimensions and datasets for test 

660 registry.insertDimensionData( 

661 "instrument", 

662 dict(name="DummyCam", visit_max=25, exposure_max=300, detector_max=6) 

663 ) 

664 registry.insertDimensionData( 

665 "physical_filter", 

666 dict(instrument="DummyCam", name="dummy_r", band="r"), 

667 dict(instrument="DummyCam", name="dummy_i", band="i"), 

668 ) 

669 registry.insertDimensionData( 

670 "detector", 

671 *[dict(instrument="DummyCam", id=i, full_name=str(i)) for i in range(1, 6)] 

672 ) 

673 registry.insertDimensionData( 

674 "visit_system", 

675 dict(instrument="DummyCam", id=1, name="default"), 

676 ) 

677 registry.insertDimensionData( 

678 "visit", 

679 dict(instrument="DummyCam", id=10, name="ten", physical_filter="dummy_i", visit_system=1), 

680 dict(instrument="DummyCam", id=11, name="eleven", physical_filter="dummy_r", visit_system=1), 

681 dict(instrument="DummyCam", id=20, name="twelve", physical_filter="dummy_r", visit_system=1), 

682 ) 

683 registry.insertDimensionData( 

684 "exposure", 

685 dict(instrument="DummyCam", id=100, name="100", physical_filter="dummy_i"), 

686 dict(instrument="DummyCam", id=101, name="101", physical_filter="dummy_i"), 

687 dict(instrument="DummyCam", id=110, name="110", physical_filter="dummy_r"), 

688 dict(instrument="DummyCam", id=111, name="111", physical_filter="dummy_r"), 

689 dict(instrument="DummyCam", id=200, name="200", physical_filter="dummy_r"), 

690 dict(instrument="DummyCam", id=201, name="201", physical_filter="dummy_r"), 

691 ) 

692 registry.insertDimensionData( 

693 "visit_definition", 

694 dict(instrument="DummyCam", exposure=100, visit_system=1, visit=10), 

695 dict(instrument="DummyCam", exposure=101, visit_system=1, visit=10), 

696 dict(instrument="DummyCam", exposure=110, visit_system=1, visit=11), 

697 dict(instrument="DummyCam", exposure=111, visit_system=1, visit=11), 

698 dict(instrument="DummyCam", exposure=200, visit_system=1, visit=20), 

699 dict(instrument="DummyCam", exposure=201, visit_system=1, visit=20), 

700 ) 

701 # dataset types 

702 run1 = "test1_r" 

703 run2 = "test2_r" 

704 tagged2 = "test2_t" 

705 registry.registerRun(run1) 

706 registry.registerRun(run2) 

707 registry.registerCollection(tagged2) 

708 storageClass = StorageClass("testDataset") 

709 registry.storageClasses.registerStorageClass(storageClass) 

710 rawType = DatasetType(name="RAW", 

711 dimensions=registry.dimensions.extract(("instrument", "exposure", "detector")), 

712 storageClass=storageClass) 

713 registry.registerDatasetType(rawType) 

714 calexpType = DatasetType(name="CALEXP", 

715 dimensions=registry.dimensions.extract(("instrument", "visit", "detector")), 

716 storageClass=storageClass) 

717 registry.registerDatasetType(calexpType) 

718 

719 # add pre-existing datasets 

720 for exposure in (100, 101, 110, 111): 

721 for detector in (1, 2, 3): 

722 # note that only 3 of 5 detectors have datasets 

723 dataId = dict(instrument="DummyCam", exposure=exposure, detector=detector) 

724 ref, = registry.insertDatasets(rawType, dataIds=[dataId], run=run1) 

725 # exposures 100 and 101 appear in both run1 and tagged2. 

726 # 100 has different datasets in the different collections 

727 # 101 has the same dataset in both collections. 

728 if exposure == 100: 

729 ref, = registry.insertDatasets(rawType, dataIds=[dataId], run=run2) 

730 if exposure in (100, 101): 

731 registry.associate(tagged2, [ref]) 

732 # Add pre-existing datasets to tagged2. 

733 for exposure in (200, 201): 

734 for detector in (3, 4, 5): 

735 # note that only 3 of 5 detectors have datasets 

736 dataId = dict(instrument="DummyCam", exposure=exposure, detector=detector) 

737 ref, = registry.insertDatasets(rawType, dataIds=[dataId], run=run2) 

738 registry.associate(tagged2, [ref]) 

739 

740 dimensions = DimensionGraph( 

741 registry.dimensions, 

742 dimensions=(rawType.dimensions.required | calexpType.dimensions.required) 

743 ) 

744 # Test that single dim string works as well as list of str 

745 rows = registry.queryDataIds("visit", datasets=rawType, collections=run1).expanded().toSet() 

746 rowsI = registry.queryDataIds(["visit"], datasets=rawType, collections=run1).expanded().toSet() 

747 self.assertEqual(rows, rowsI) 

748 # with empty expression 

749 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=run1).expanded().toSet() 

750 self.assertEqual(len(rows), 4*3) # 4 exposures times 3 detectors 

751 for dataId in rows: 

752 self.assertCountEqual(dataId.keys(), ("instrument", "detector", "exposure", "visit")) 

753 packer1 = registry.dimensions.makePacker("visit_detector", dataId) 

754 packer2 = registry.dimensions.makePacker("exposure_detector", dataId) 

755 self.assertEqual(packer1.unpack(packer1.pack(dataId)), 

756 DataCoordinate.standardize(dataId, graph=packer1.dimensions)) 

757 self.assertEqual(packer2.unpack(packer2.pack(dataId)), 

758 DataCoordinate.standardize(dataId, graph=packer2.dimensions)) 

759 self.assertNotEqual(packer1.pack(dataId), packer2.pack(dataId)) 

760 self.assertCountEqual(set(dataId["exposure"] for dataId in rows), 

761 (100, 101, 110, 111)) 

762 self.assertCountEqual(set(dataId["visit"] for dataId in rows), (10, 11)) 

763 self.assertCountEqual(set(dataId["detector"] for dataId in rows), (1, 2, 3)) 

764 

765 # second collection 

766 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=tagged2).toSet() 

767 self.assertEqual(len(rows), 4*3) # 4 exposures times 3 detectors 

768 for dataId in rows: 

769 self.assertCountEqual(dataId.keys(), ("instrument", "detector", "exposure", "visit")) 

770 self.assertCountEqual(set(dataId["exposure"] for dataId in rows), 

771 (100, 101, 200, 201)) 

772 self.assertCountEqual(set(dataId["visit"] for dataId in rows), (10, 20)) 

773 self.assertCountEqual(set(dataId["detector"] for dataId in rows), (1, 2, 3, 4, 5)) 

774 

775 # with two input datasets 

776 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=[run1, tagged2]).toSet() 

777 self.assertEqual(len(set(rows)), 6*3) # 6 exposures times 3 detectors; set needed to de-dupe 

778 for dataId in rows: 

779 self.assertCountEqual(dataId.keys(), ("instrument", "detector", "exposure", "visit")) 

780 self.assertCountEqual(set(dataId["exposure"] for dataId in rows), 

781 (100, 101, 110, 111, 200, 201)) 

782 self.assertCountEqual(set(dataId["visit"] for dataId in rows), (10, 11, 20)) 

783 self.assertCountEqual(set(dataId["detector"] for dataId in rows), (1, 2, 3, 4, 5)) 

784 

785 # limit to single visit 

786 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=run1, 

787 where="visit = 10").toSet() 

788 self.assertEqual(len(rows), 2*3) # 2 exposures times 3 detectors 

789 self.assertCountEqual(set(dataId["exposure"] for dataId in rows), (100, 101)) 

790 self.assertCountEqual(set(dataId["visit"] for dataId in rows), (10,)) 

791 self.assertCountEqual(set(dataId["detector"] for dataId in rows), (1, 2, 3)) 

792 

793 # more limiting expression, using link names instead of Table.column 

794 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=run1, 

795 where="visit = 10 and detector > 1").toSet() 

796 self.assertEqual(len(rows), 2*2) # 2 exposures times 2 detectors 

797 self.assertCountEqual(set(dataId["exposure"] for dataId in rows), (100, 101)) 

798 self.assertCountEqual(set(dataId["visit"] for dataId in rows), (10,)) 

799 self.assertCountEqual(set(dataId["detector"] for dataId in rows), (2, 3)) 

800 

801 # expression excludes everything 

802 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=run1, 

803 where="visit > 1000").toSet() 

804 self.assertEqual(len(rows), 0) 

805 

806 # Selecting by physical_filter, this is not in the dimensions, but it 

807 # is a part of the full expression so it should work too. 

808 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=run1, 

809 where="physical_filter = 'dummy_r'").toSet() 

810 self.assertEqual(len(rows), 2*3) # 2 exposures times 3 detectors 

811 self.assertCountEqual(set(dataId["exposure"] for dataId in rows), (110, 111)) 

812 self.assertCountEqual(set(dataId["visit"] for dataId in rows), (11,)) 

813 self.assertCountEqual(set(dataId["detector"] for dataId in rows), (1, 2, 3)) 

814 

815 def testSkyMapDimensions(self): 

816 """Tests involving only skymap dimensions, no joins to instrument.""" 

817 registry = self.makeRegistry() 

818 

819 # need a bunch of dimensions and datasets for test, we want 

820 # "band" in the test so also have to add physical_filter 

821 # dimensions 

822 registry.insertDimensionData( 

823 "instrument", 

824 dict(instrument="DummyCam") 

825 ) 

826 registry.insertDimensionData( 

827 "physical_filter", 

828 dict(instrument="DummyCam", name="dummy_r", band="r"), 

829 dict(instrument="DummyCam", name="dummy_i", band="i"), 

830 ) 

831 registry.insertDimensionData( 

832 "skymap", 

833 dict(name="DummyMap", hash="sha!".encode("utf8")) 

834 ) 

835 for tract in range(10): 

836 registry.insertDimensionData("tract", dict(skymap="DummyMap", id=tract)) 

837 registry.insertDimensionData( 

838 "patch", 

839 *[dict(skymap="DummyMap", tract=tract, id=patch, cell_x=0, cell_y=0) 

840 for patch in range(10)] 

841 ) 

842 

843 # dataset types 

844 run = "test" 

845 registry.registerRun(run) 

846 storageClass = StorageClass("testDataset") 

847 registry.storageClasses.registerStorageClass(storageClass) 

848 calexpType = DatasetType(name="deepCoadd_calexp", 

849 dimensions=registry.dimensions.extract(("skymap", "tract", "patch", 

850 "band")), 

851 storageClass=storageClass) 

852 registry.registerDatasetType(calexpType) 

853 mergeType = DatasetType(name="deepCoadd_mergeDet", 

854 dimensions=registry.dimensions.extract(("skymap", "tract", "patch")), 

855 storageClass=storageClass) 

856 registry.registerDatasetType(mergeType) 

857 measType = DatasetType(name="deepCoadd_meas", 

858 dimensions=registry.dimensions.extract(("skymap", "tract", "patch", 

859 "band")), 

860 storageClass=storageClass) 

861 registry.registerDatasetType(measType) 

862 

863 dimensions = DimensionGraph( 

864 registry.dimensions, 

865 dimensions=(calexpType.dimensions.required | mergeType.dimensions.required 

866 | measType.dimensions.required) 

867 ) 

868 

869 # add pre-existing datasets 

870 for tract in (1, 3, 5): 

871 for patch in (2, 4, 6, 7): 

872 dataId = dict(skymap="DummyMap", tract=tract, patch=patch) 

873 registry.insertDatasets(mergeType, dataIds=[dataId], run=run) 

874 for aFilter in ("i", "r"): 

875 dataId = dict(skymap="DummyMap", tract=tract, patch=patch, band=aFilter) 

876 registry.insertDatasets(calexpType, dataIds=[dataId], run=run) 

877 

878 # with empty expression 

879 rows = registry.queryDataIds(dimensions, 

880 datasets=[calexpType, mergeType], collections=run).toSet() 

881 self.assertEqual(len(rows), 3*4*2) # 4 tracts x 4 patches x 2 filters 

882 for dataId in rows: 

883 self.assertCountEqual(dataId.keys(), ("skymap", "tract", "patch", "band")) 

884 self.assertCountEqual(set(dataId["tract"] for dataId in rows), (1, 3, 5)) 

885 self.assertCountEqual(set(dataId["patch"] for dataId in rows), (2, 4, 6, 7)) 

886 self.assertCountEqual(set(dataId["band"] for dataId in rows), ("i", "r")) 

887 

888 # limit to 2 tracts and 2 patches 

889 rows = registry.queryDataIds(dimensions, 

890 datasets=[calexpType, mergeType], collections=run, 

891 where="tract IN (1, 5) AND patch IN (2, 7)").toSet() 

892 self.assertEqual(len(rows), 2*2*2) # 2 tracts x 2 patches x 2 filters 

893 self.assertCountEqual(set(dataId["tract"] for dataId in rows), (1, 5)) 

894 self.assertCountEqual(set(dataId["patch"] for dataId in rows), (2, 7)) 

895 self.assertCountEqual(set(dataId["band"] for dataId in rows), ("i", "r")) 

896 

897 # limit to single filter 

898 rows = registry.queryDataIds(dimensions, 

899 datasets=[calexpType, mergeType], collections=run, 

900 where="band = 'i'").toSet() 

901 self.assertEqual(len(rows), 3*4*1) # 4 tracts x 4 patches x 2 filters 

902 self.assertCountEqual(set(dataId["tract"] for dataId in rows), (1, 3, 5)) 

903 self.assertCountEqual(set(dataId["patch"] for dataId in rows), (2, 4, 6, 7)) 

904 self.assertCountEqual(set(dataId["band"] for dataId in rows), ("i",)) 

905 

906 # expression excludes everything, specifying non-existing skymap is 

907 # not a fatal error, it's operator error 

908 rows = registry.queryDataIds(dimensions, 

909 datasets=[calexpType, mergeType], collections=run, 

910 where="skymap = 'Mars'").toSet() 

911 self.assertEqual(len(rows), 0) 

912 

913 def testSpatialMatch(self): 

914 """Test involving spatial match using join tables. 

915 

916 Note that realistic test needs a reasonably-defined skypix and regions 

917 in registry tables which is hard to implement in this simple test. 

918 So we do not actually fill registry with any data and all queries will 

919 return empty result, but this is still useful for coverage of the code 

920 that generates query. 

921 """ 

922 registry = self.makeRegistry() 

923 

924 # dataset types 

925 collection = "test" 

926 registry.registerRun(name=collection) 

927 storageClass = StorageClass("testDataset") 

928 registry.storageClasses.registerStorageClass(storageClass) 

929 

930 calexpType = DatasetType(name="CALEXP", 

931 dimensions=registry.dimensions.extract(("instrument", "visit", "detector")), 

932 storageClass=storageClass) 

933 registry.registerDatasetType(calexpType) 

934 

935 coaddType = DatasetType(name="deepCoadd_calexp", 

936 dimensions=registry.dimensions.extract(("skymap", "tract", "patch", 

937 "band")), 

938 storageClass=storageClass) 

939 registry.registerDatasetType(coaddType) 

940 

941 dimensions = DimensionGraph( 

942 registry.dimensions, 

943 dimensions=(calexpType.dimensions.required | coaddType.dimensions.required) 

944 ) 

945 

946 # without data this should run OK but return empty set 

947 rows = registry.queryDataIds(dimensions, datasets=calexpType, collections=collection).toSet() 

948 self.assertEqual(len(rows), 0) 

949 

950 def testAbstractQuery(self): 

951 """Test that we can run a query that just lists the known 

952 bands. This is tricky because band is 

953 backed by a query against physical_filter. 

954 """ 

955 registry = self.makeRegistry() 

956 registry.insertDimensionData("instrument", dict(name="DummyCam")) 

957 registry.insertDimensionData( 

958 "physical_filter", 

959 dict(instrument="DummyCam", name="dummy_i", band="i"), 

960 dict(instrument="DummyCam", name="dummy_i2", band="i"), 

961 dict(instrument="DummyCam", name="dummy_r", band="r"), 

962 ) 

963 rows = registry.queryDataIds(["band"]).toSet() 

964 self.assertCountEqual( 

965 rows, 

966 [DataCoordinate.standardize(band="i", universe=registry.dimensions), 

967 DataCoordinate.standardize(band="r", universe=registry.dimensions)] 

968 ) 

969 

970 def testAttributeManager(self): 

971 """Test basic functionality of attribute manager. 

972 """ 

973 # number of attributes with schema versions in a fresh database, 

974 # 6 managers with 3 records per manager 

975 VERSION_COUNT = 6 * 3 

976 

977 registry = self.makeRegistry() 

978 attributes = registry._attributes 

979 

980 # check what get() returns for non-existing key 

981 self.assertIsNone(attributes.get("attr")) 

982 self.assertEqual(attributes.get("attr", ""), "") 

983 self.assertEqual(attributes.get("attr", "Value"), "Value") 

984 self.assertEqual(len(list(attributes.items())), VERSION_COUNT) 

985 

986 # cannot store empty key or value 

987 with self.assertRaises(ValueError): 

988 attributes.set("", "value") 

989 with self.assertRaises(ValueError): 

990 attributes.set("attr", "") 

991 

992 # set value of non-existing key 

993 attributes.set("attr", "value") 

994 self.assertEqual(len(list(attributes.items())), VERSION_COUNT + 1) 

995 self.assertEqual(attributes.get("attr"), "value") 

996 

997 # update value of existing key 

998 with self.assertRaises(ButlerAttributeExistsError): 

999 attributes.set("attr", "value2") 

1000 

1001 attributes.set("attr", "value2", force=True) 

1002 self.assertEqual(len(list(attributes.items())), VERSION_COUNT + 1) 

1003 self.assertEqual(attributes.get("attr"), "value2") 

1004 

1005 # delete existing key 

1006 self.assertTrue(attributes.delete("attr")) 

1007 self.assertEqual(len(list(attributes.items())), VERSION_COUNT) 

1008 

1009 # delete non-existing key 

1010 self.assertFalse(attributes.delete("non-attr")) 

1011 

1012 # store bunch of keys and get the list back 

1013 data = [ 

1014 ("version.core", "1.2.3"), 

1015 ("version.dimensions", "3.2.1"), 

1016 ("config.managers.opaque", "ByNameOpaqueTableStorageManager"), 

1017 ] 

1018 for key, value in data: 

1019 attributes.set(key, value) 

1020 items = dict(attributes.items()) 

1021 for key, value in data: 

1022 self.assertEqual(items[key], value) 

1023 

1024 def testQueryDatasetsDeduplication(self): 

1025 """Test that the findFirst option to queryDatasets selects datasets 

1026 from collections in the order given". 

1027 """ 

1028 registry = self.makeRegistry() 

1029 self.loadData(registry, "base.yaml") 

1030 self.loadData(registry, "datasets.yaml") 

1031 self.assertCountEqual( 

1032 list(registry.queryDatasets("bias", collections=["imported_g", "imported_r"])), 

1033 [ 

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

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

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

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

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

1039 registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"), 

1040 ] 

1041 ) 

1042 self.assertCountEqual( 

1043 list(registry.queryDatasets("bias", collections=["imported_g", "imported_r"], 

1044 findFirst=True)), 

1045 [ 

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

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

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

1049 registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"), 

1050 ] 

1051 ) 

1052 self.assertCountEqual( 

1053 list(registry.queryDatasets("bias", collections=["imported_r", "imported_g"], 

1054 findFirst=True)), 

1055 [ 

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

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

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

1059 registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"), 

1060 ] 

1061 ) 

1062 

1063 def testQueryResults(self): 

1064 """Test querying for data IDs and then manipulating the QueryResults 

1065 object returned to perform other queries. 

1066 """ 

1067 registry = self.makeRegistry() 

1068 self.loadData(registry, "base.yaml") 

1069 self.loadData(registry, "datasets.yaml") 

1070 bias = registry.getDatasetType("bias") 

1071 flat = registry.getDatasetType("flat") 

1072 # Obtain expected results from methods other than those we're testing 

1073 # here. That includes: 

1074 # - the dimensions of the data IDs we want to query: 

1075 expectedGraph = DimensionGraph(registry.dimensions, names=["detector", "physical_filter"]) 

1076 # - the dimensions of some other data IDs we'll extract from that: 

1077 expectedSubsetGraph = DimensionGraph(registry.dimensions, names=["detector"]) 

1078 # - the data IDs we expect to obtain from the first queries: 

1079 expectedDataIds = DataCoordinateSet( 

1080 { 

1081 DataCoordinate.standardize(instrument="Cam1", detector=d, physical_filter=p, 

1082 universe=registry.dimensions) 

1083 for d, p in itertools.product({1, 2, 3}, {"Cam1-G", "Cam1-R1", "Cam1-R2"}) 

1084 }, 

1085 graph=expectedGraph, 

1086 hasFull=False, 

1087 hasRecords=False, 

1088 ) 

1089 # - the flat datasets we expect to find from those data IDs, in just 

1090 # one collection (so deduplication is irrelevant): 

1091 expectedFlats = [ 

1092 registry.findDataset(flat, instrument="Cam1", detector=1, physical_filter="Cam1-R1", 

1093 collections="imported_r"), 

1094 registry.findDataset(flat, instrument="Cam1", detector=2, physical_filter="Cam1-R1", 

1095 collections="imported_r"), 

1096 registry.findDataset(flat, instrument="Cam1", detector=3, physical_filter="Cam1-R2", 

1097 collections="imported_r"), 

1098 ] 

1099 # - the data IDs we expect to extract from that: 

1100 expectedSubsetDataIds = expectedDataIds.subset(expectedSubsetGraph) 

1101 # - the bias datasets we expect to find from those data IDs, after we 

1102 # subset-out the physical_filter dimension, both with duplicates: 

1103 expectedAllBiases = [ 

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

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

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

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

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

1109 ] 

1110 # - ...and without duplicates: 

1111 expectedDeduplicatedBiases = [ 

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

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

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

1115 ] 

1116 # Test against those expected results, using a "lazy" query for the 

1117 # data IDs (which re-executes that query each time we use it to do 

1118 # something new). 

1119 dataIds = registry.queryDataIds( 

1120 ["detector", "physical_filter"], 

1121 where="detector.purpose = 'SCIENCE'", # this rejects detector=4 

1122 ) 

1123 self.assertEqual(dataIds.graph, expectedGraph) 

1124 self.assertEqual(dataIds.toSet(), expectedDataIds) 

1125 self.assertCountEqual( 

1126 list( 

1127 dataIds.findDatasets( 

1128 flat, 

1129 collections=["imported_r"], 

1130 ) 

1131 ), 

1132 expectedFlats, 

1133 ) 

1134 subsetDataIds = dataIds.subset(expectedSubsetGraph, unique=True) 

1135 self.assertEqual(subsetDataIds.graph, expectedSubsetGraph) 

1136 self.assertEqual(subsetDataIds.toSet(), expectedSubsetDataIds) 

1137 self.assertCountEqual( 

1138 list( 

1139 subsetDataIds.findDatasets( 

1140 bias, 

1141 collections=["imported_r", "imported_g"], 

1142 findFirst=False 

1143 ) 

1144 ), 

1145 expectedAllBiases 

1146 ) 

1147 self.assertCountEqual( 

1148 list( 

1149 subsetDataIds.findDatasets( 

1150 bias, 

1151 collections=["imported_r", "imported_g"], 

1152 findFirst=True 

1153 ) 

1154 ), expectedDeduplicatedBiases 

1155 ) 

1156 # Materialize the bias dataset queries (only) by putting the results 

1157 # into temporary tables, then repeat those tests. 

1158 with subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], 

1159 findFirst=False).materialize() as biases: 

1160 self.assertCountEqual(list(biases), expectedAllBiases) 

1161 with subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], 

1162 findFirst=True).materialize() as biases: 

1163 self.assertCountEqual(list(biases), expectedDeduplicatedBiases) 

1164 # Materialize the data ID subset query, but not the dataset queries. 

1165 with subsetDataIds.materialize() as subsetDataIds: 

1166 self.assertEqual(subsetDataIds.graph, expectedSubsetGraph) 

1167 self.assertEqual(subsetDataIds.toSet(), expectedSubsetDataIds) 

1168 self.assertCountEqual( 

1169 list( 

1170 subsetDataIds.findDatasets( 

1171 bias, 

1172 collections=["imported_r", "imported_g"], 

1173 findFirst=False 

1174 ) 

1175 ), 

1176 expectedAllBiases 

1177 ) 

1178 self.assertCountEqual( 

1179 list( 

1180 subsetDataIds.findDatasets( 

1181 bias, 

1182 collections=["imported_r", "imported_g"], 

1183 findFirst=True 

1184 ) 

1185 ), expectedDeduplicatedBiases 

1186 ) 

1187 # Materialize the dataset queries, too. 

1188 with subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], 

1189 findFirst=False).materialize() as biases: 

1190 self.assertCountEqual(list(biases), expectedAllBiases) 

1191 with subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], 

1192 findFirst=True).materialize() as biases: 

1193 self.assertCountEqual(list(biases), expectedDeduplicatedBiases) 

1194 # Materialize the original query, but none of the follow-up queries. 

1195 with dataIds.materialize() as dataIds: 

1196 self.assertEqual(dataIds.graph, expectedGraph) 

1197 self.assertEqual(dataIds.toSet(), expectedDataIds) 

1198 self.assertCountEqual( 

1199 list( 

1200 dataIds.findDatasets( 

1201 flat, 

1202 collections=["imported_r"], 

1203 ) 

1204 ), 

1205 expectedFlats, 

1206 ) 

1207 subsetDataIds = dataIds.subset(expectedSubsetGraph, unique=True) 

1208 self.assertEqual(subsetDataIds.graph, expectedSubsetGraph) 

1209 self.assertEqual(subsetDataIds.toSet(), expectedSubsetDataIds) 

1210 self.assertCountEqual( 

1211 list( 

1212 subsetDataIds.findDatasets( 

1213 bias, 

1214 collections=["imported_r", "imported_g"], 

1215 findFirst=False 

1216 ) 

1217 ), 

1218 expectedAllBiases 

1219 ) 

1220 self.assertCountEqual( 

1221 list( 

1222 subsetDataIds.findDatasets( 

1223 bias, 

1224 collections=["imported_r", "imported_g"], 

1225 findFirst=True 

1226 ) 

1227 ), expectedDeduplicatedBiases 

1228 ) 

1229 # Materialize just the bias dataset queries. 

1230 with subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], 

1231 findFirst=False).materialize() as biases: 

1232 self.assertCountEqual(list(biases), expectedAllBiases) 

1233 with subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], 

1234 findFirst=True).materialize() as biases: 

1235 self.assertCountEqual(list(biases), expectedDeduplicatedBiases) 

1236 # Materialize the subset data ID query, but not the dataset 

1237 # queries. 

1238 with subsetDataIds.materialize() as subsetDataIds: 

1239 self.assertEqual(subsetDataIds.graph, expectedSubsetGraph) 

1240 self.assertEqual(subsetDataIds.toSet(), expectedSubsetDataIds) 

1241 self.assertCountEqual( 

1242 list( 

1243 subsetDataIds.findDatasets( 

1244 bias, 

1245 collections=["imported_r", "imported_g"], 

1246 findFirst=False 

1247 ) 

1248 ), 

1249 expectedAllBiases 

1250 ) 

1251 self.assertCountEqual( 

1252 list( 

1253 subsetDataIds.findDatasets( 

1254 bias, 

1255 collections=["imported_r", "imported_g"], 

1256 findFirst=True 

1257 ) 

1258 ), expectedDeduplicatedBiases 

1259 ) 

1260 # Materialize the bias dataset queries, too, so now we're 

1261 # materializing every single step. 

1262 with subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], 

1263 findFirst=False).materialize() as biases: 

1264 self.assertCountEqual(list(biases), expectedAllBiases) 

1265 with subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], 

1266 findFirst=True).materialize() as biases: 

1267 self.assertCountEqual(list(biases), expectedDeduplicatedBiases) 

1268 

1269 def testEmptyDimensionsQueries(self): 

1270 """Test Query and QueryResults objects in the case where there are no 

1271 dimensions. 

1272 """ 

1273 # Set up test data: one dataset type, two runs, one dataset in each. 

1274 registry = self.makeRegistry() 

1275 self.loadData(registry, "base.yaml") 

1276 schema = DatasetType("schema", dimensions=registry.dimensions.empty, storageClass="Catalog") 

1277 registry.registerDatasetType(schema) 

1278 dataId = DataCoordinate.makeEmpty(registry.dimensions) 

1279 run1 = "run1" 

1280 run2 = "run2" 

1281 registry.registerRun(run1) 

1282 registry.registerRun(run2) 

1283 (dataset1,) = registry.insertDatasets(schema, dataIds=[dataId], run=run1) 

1284 (dataset2,) = registry.insertDatasets(schema, dataIds=[dataId], run=run2) 

1285 # Query directly for both of the datasets, and each one, one at a time. 

1286 self.assertCountEqual( 

1287 list(registry.queryDatasets(schema, collections=[run1, run2], findFirst=False)), 

1288 [dataset1, dataset2] 

1289 ) 

1290 self.assertEqual( 

1291 list(registry.queryDatasets(schema, collections=[run1, run2], findFirst=True)), 

1292 [dataset1], 

1293 ) 

1294 self.assertEqual( 

1295 list(registry.queryDatasets(schema, collections=[run2, run1], findFirst=True)), 

1296 [dataset2], 

1297 ) 

1298 # Query for data IDs with no dimensions. 

1299 dataIds = registry.queryDataIds([]) 

1300 self.assertEqual( 

1301 dataIds.toSequence(), 

1302 DataCoordinateSequence([dataId], registry.dimensions.empty) 

1303 ) 

1304 # Use queried data IDs to find the datasets. 

1305 self.assertCountEqual( 

1306 list(dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False)), 

1307 [dataset1, dataset2], 

1308 ) 

1309 self.assertEqual( 

1310 list(dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True)), 

1311 [dataset1], 

1312 ) 

1313 self.assertEqual( 

1314 list(dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True)), 

1315 [dataset2], 

1316 ) 

1317 # Now materialize the data ID query results and repeat those tests. 

1318 with dataIds.materialize() as dataIds: 

1319 self.assertEqual( 

1320 dataIds.toSequence(), 

1321 DataCoordinateSequence([dataId], registry.dimensions.empty) 

1322 ) 

1323 self.assertCountEqual( 

1324 list(dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False)), 

1325 [dataset1, dataset2], 

1326 ) 

1327 self.assertEqual( 

1328 list(dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True)), 

1329 [dataset1], 

1330 ) 

1331 self.assertEqual( 

1332 list(dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True)), 

1333 [dataset2], 

1334 ) 

1335 # Query for non-empty data IDs, then subset that to get the empty one. 

1336 # Repeat the above tests starting from that. 

1337 dataIds = registry.queryDataIds(["instrument"]).subset(registry.dimensions.empty, unique=True) 

1338 self.assertEqual( 

1339 dataIds.toSequence(), 

1340 DataCoordinateSequence([dataId], registry.dimensions.empty) 

1341 ) 

1342 self.assertCountEqual( 

1343 list(dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False)), 

1344 [dataset1, dataset2], 

1345 ) 

1346 self.assertEqual( 

1347 list(dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True)), 

1348 [dataset1], 

1349 ) 

1350 self.assertEqual( 

1351 list(dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True)), 

1352 [dataset2], 

1353 ) 

1354 with dataIds.materialize() as dataIds: 

1355 self.assertEqual( 

1356 dataIds.toSequence(), 

1357 DataCoordinateSequence([dataId], registry.dimensions.empty) 

1358 ) 

1359 self.assertCountEqual( 

1360 list(dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False)), 

1361 [dataset1, dataset2], 

1362 ) 

1363 self.assertEqual( 

1364 list(dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True)), 

1365 [dataset1], 

1366 ) 

1367 self.assertEqual( 

1368 list(dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True)), 

1369 [dataset2], 

1370 ) 

1371 # Query for non-empty data IDs, then materialize, then subset to get 

1372 # the empty one. Repeat again. 

1373 with registry.queryDataIds(["instrument"]).materialize() as nonEmptyDataIds: 

1374 dataIds = nonEmptyDataIds.subset(registry.dimensions.empty, unique=True) 

1375 self.assertEqual( 

1376 dataIds.toSequence(), 

1377 DataCoordinateSequence([dataId], registry.dimensions.empty) 

1378 ) 

1379 self.assertCountEqual( 

1380 list(dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False)), 

1381 [dataset1, dataset2], 

1382 ) 

1383 self.assertEqual( 

1384 list(dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True)), 

1385 [dataset1], 

1386 ) 

1387 self.assertEqual( 

1388 list(dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True)), 

1389 [dataset2], 

1390 ) 

1391 with dataIds.materialize() as dataIds: 

1392 self.assertEqual( 

1393 dataIds.toSequence(), 

1394 DataCoordinateSequence([dataId], registry.dimensions.empty) 

1395 ) 

1396 self.assertCountEqual( 

1397 list(dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False)), 

1398 [dataset1, dataset2], 

1399 ) 

1400 self.assertEqual( 

1401 list(dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True)), 

1402 [dataset1], 

1403 ) 

1404 self.assertEqual( 

1405 list(dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True)), 

1406 [dataset2], 

1407 ) 

1408 

1409 def testCalibrationCollections(self): 

1410 """Test operations on `~CollectionType.CALIBRATION` collections, 

1411 including `Registry.certify`, `Registry.decertify`, and 

1412 `Registry.findDataset`. 

1413 """ 

1414 # Setup - make a Registry, fill it with some datasets in 

1415 # non-calibration collections. 

1416 registry = self.makeRegistry() 

1417 self.loadData(registry, "base.yaml") 

1418 self.loadData(registry, "datasets.yaml") 

1419 # Set up some timestamps. 

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

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

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

1423 t4 = astropy.time.Time('2020-01-01T04:00:00', format="isot", scale="tai") 

1424 t5 = astropy.time.Time('2020-01-01T05:00:00', format="isot", scale="tai") 

1425 allTimespans = [ 

1426 Timespan(a, b) for a, b in itertools.combinations([None, t1, t2, t3, t4, t5, None], r=2) 

1427 ] 

1428 # Get references to some datasets. 

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

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

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

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

1433 # Register the main calibration collection we'll be working with. 

1434 collection = "Cam1/calibs/default" 

1435 registry.registerCollection(collection, type=CollectionType.CALIBRATION) 

1436 # Cannot associate into a calibration collection (no timespan). 

1437 with self.assertRaises(TypeError): 

1438 registry.associate(collection, [bias2a]) 

1439 # Certify 2a dataset with [t2, t4) validity. 

1440 registry.certify(collection, [bias2a], Timespan(begin=t2, end=t4)) 

1441 # We should not be able to certify 2b with anything overlapping that 

1442 # window. 

1443 with self.assertRaises(ConflictingDefinitionError): 

1444 registry.certify(collection, [bias2b], Timespan(begin=None, end=t3)) 

1445 with self.assertRaises(ConflictingDefinitionError): 

1446 registry.certify(collection, [bias2b], Timespan(begin=None, end=t5)) 

1447 with self.assertRaises(ConflictingDefinitionError): 

1448 registry.certify(collection, [bias2b], Timespan(begin=t1, end=t3)) 

1449 with self.assertRaises(ConflictingDefinitionError): 

1450 registry.certify(collection, [bias2b], Timespan(begin=t1, end=t5)) 

1451 with self.assertRaises(ConflictingDefinitionError): 

1452 registry.certify(collection, [bias2b], Timespan(begin=t1, end=None)) 

1453 with self.assertRaises(ConflictingDefinitionError): 

1454 registry.certify(collection, [bias2b], Timespan(begin=t2, end=t3)) 

1455 with self.assertRaises(ConflictingDefinitionError): 

1456 registry.certify(collection, [bias2b], Timespan(begin=t2, end=t5)) 

1457 with self.assertRaises(ConflictingDefinitionError): 

1458 registry.certify(collection, [bias2b], Timespan(begin=t2, end=None)) 

1459 # We should be able to certify 3a with a range overlapping that window, 

1460 # because it's for a different detector. 

1461 # We'll certify 3a over [t1, t3). 

1462 registry.certify(collection, [bias3a], Timespan(begin=t1, end=t3)) 

1463 # Now we'll certify 2b and 3b together over [t4, ∞). 

1464 registry.certify(collection, [bias2b, bias3b], Timespan(begin=t4, end=None)) 

1465 

1466 # Fetch all associations and check that they are what we expect. 

1467 self.assertCountEqual( 

1468 list( 

1469 registry.queryDatasetAssociations( 

1470 "bias", 

1471 collections=[collection, "imported_g", "imported_r"], 

1472 ) 

1473 ), 

1474 [ 

1475 DatasetAssociation( 

1476 ref=registry.findDataset("bias", instrument="Cam1", detector=1, collections="imported_g"), 

1477 collection="imported_g", 

1478 timespan=None, 

1479 ), 

1480 DatasetAssociation( 

1481 ref=registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"), 

1482 collection="imported_r", 

1483 timespan=None, 

1484 ), 

1485 DatasetAssociation(ref=bias2a, collection="imported_g", timespan=None), 

1486 DatasetAssociation(ref=bias3a, collection="imported_g", timespan=None), 

1487 DatasetAssociation(ref=bias2b, collection="imported_r", timespan=None), 

1488 DatasetAssociation(ref=bias3b, collection="imported_r", timespan=None), 

1489 DatasetAssociation(ref=bias2a, collection=collection, timespan=Timespan(begin=t2, end=t4)), 

1490 DatasetAssociation(ref=bias3a, collection=collection, timespan=Timespan(begin=t1, end=t3)), 

1491 DatasetAssociation(ref=bias2b, collection=collection, timespan=Timespan(begin=t4, end=None)), 

1492 DatasetAssociation(ref=bias3b, collection=collection, timespan=Timespan(begin=t4, end=None)), 

1493 ] 

1494 ) 

1495 

1496 class Ambiguous: 

1497 """Tag class to denote lookups that are expected to be ambiguous. 

1498 """ 

1499 pass 

1500 

1501 def assertLookup(detector: int, timespan: Timespan, 

1502 expected: Optional[Union[DatasetRef, Type[Ambiguous]]]) -> None: 

1503 """Local function that asserts that a bias lookup returns the given 

1504 expected result. 

1505 """ 

1506 if expected is Ambiguous: 

1507 with self.assertRaises(RuntimeError): 

1508 registry.findDataset("bias", collections=collection, instrument="Cam1", 

1509 detector=detector, timespan=timespan) 

1510 else: 

1511 self.assertEqual( 

1512 expected, 

1513 registry.findDataset("bias", collections=collection, instrument="Cam1", 

1514 detector=detector, timespan=timespan) 

1515 ) 

1516 

1517 # Systematically test lookups against expected results. 

1518 assertLookup(detector=2, timespan=Timespan(None, t1), expected=None) 

1519 assertLookup(detector=2, timespan=Timespan(None, t2), expected=None) 

1520 assertLookup(detector=2, timespan=Timespan(None, t3), expected=bias2a) 

1521 assertLookup(detector=2, timespan=Timespan(None, t4), expected=bias2a) 

1522 assertLookup(detector=2, timespan=Timespan(None, t5), expected=Ambiguous) 

1523 assertLookup(detector=2, timespan=Timespan(None, None), expected=Ambiguous) 

1524 assertLookup(detector=2, timespan=Timespan(t1, t2), expected=None) 

1525 assertLookup(detector=2, timespan=Timespan(t1, t3), expected=bias2a) 

1526 assertLookup(detector=2, timespan=Timespan(t1, t4), expected=bias2a) 

1527 assertLookup(detector=2, timespan=Timespan(t1, t5), expected=Ambiguous) 

1528 assertLookup(detector=2, timespan=Timespan(t1, None), expected=Ambiguous) 

1529 assertLookup(detector=2, timespan=Timespan(t2, t3), expected=bias2a) 

1530 assertLookup(detector=2, timespan=Timespan(t2, t4), expected=bias2a) 

1531 assertLookup(detector=2, timespan=Timespan(t2, t5), expected=Ambiguous) 

1532 assertLookup(detector=2, timespan=Timespan(t2, None), expected=Ambiguous) 

1533 assertLookup(detector=2, timespan=Timespan(t3, t4), expected=bias2a) 

1534 assertLookup(detector=2, timespan=Timespan(t3, t5), expected=Ambiguous) 

1535 assertLookup(detector=2, timespan=Timespan(t3, None), expected=Ambiguous) 

1536 assertLookup(detector=2, timespan=Timespan(t4, t5), expected=bias2b) 

1537 assertLookup(detector=2, timespan=Timespan(t4, None), expected=bias2b) 

1538 assertLookup(detector=2, timespan=Timespan(t5, None), expected=bias2b) 

1539 assertLookup(detector=3, timespan=Timespan(None, t1), expected=None) 

1540 assertLookup(detector=3, timespan=Timespan(None, t2), expected=bias3a) 

1541 assertLookup(detector=3, timespan=Timespan(None, t3), expected=bias3a) 

1542 assertLookup(detector=3, timespan=Timespan(None, t4), expected=bias3a) 

1543 assertLookup(detector=3, timespan=Timespan(None, t5), expected=Ambiguous) 

1544 assertLookup(detector=3, timespan=Timespan(None, None), expected=Ambiguous) 

1545 assertLookup(detector=3, timespan=Timespan(t1, t2), expected=bias3a) 

1546 assertLookup(detector=3, timespan=Timespan(t1, t3), expected=bias3a) 

1547 assertLookup(detector=3, timespan=Timespan(t1, t4), expected=bias3a) 

1548 assertLookup(detector=3, timespan=Timespan(t1, t5), expected=Ambiguous) 

1549 assertLookup(detector=3, timespan=Timespan(t1, None), expected=Ambiguous) 

1550 assertLookup(detector=3, timespan=Timespan(t2, t3), expected=bias3a) 

1551 assertLookup(detector=3, timespan=Timespan(t2, t4), expected=bias3a) 

1552 assertLookup(detector=3, timespan=Timespan(t2, t5), expected=Ambiguous) 

1553 assertLookup(detector=3, timespan=Timespan(t2, None), expected=Ambiguous) 

1554 assertLookup(detector=3, timespan=Timespan(t3, t4), expected=None) 

1555 assertLookup(detector=3, timespan=Timespan(t3, t5), expected=bias3b) 

1556 assertLookup(detector=3, timespan=Timespan(t3, None), expected=bias3b) 

1557 assertLookup(detector=3, timespan=Timespan(t4, t5), expected=bias3b) 

1558 assertLookup(detector=3, timespan=Timespan(t4, None), expected=bias3b) 

1559 assertLookup(detector=3, timespan=Timespan(t5, None), expected=bias3b) 

1560 

1561 # Decertify [t3, t5) for all data IDs, and do test lookups again. 

1562 # This should truncate bias2a to [t2, t3), leave bias3a unchanged at 

1563 # [t1, t3), and truncate bias2b and bias3b to [t5, ∞). 

1564 registry.decertify(collection=collection, datasetType="bias", timespan=Timespan(t3, t5)) 

1565 assertLookup(detector=2, timespan=Timespan(None, t1), expected=None) 

1566 assertLookup(detector=2, timespan=Timespan(None, t2), expected=None) 

1567 assertLookup(detector=2, timespan=Timespan(None, t3), expected=bias2a) 

1568 assertLookup(detector=2, timespan=Timespan(None, t4), expected=bias2a) 

1569 assertLookup(detector=2, timespan=Timespan(None, t5), expected=bias2a) 

1570 assertLookup(detector=2, timespan=Timespan(None, None), expected=Ambiguous) 

1571 assertLookup(detector=2, timespan=Timespan(t1, t2), expected=None) 

1572 assertLookup(detector=2, timespan=Timespan(t1, t3), expected=bias2a) 

1573 assertLookup(detector=2, timespan=Timespan(t1, t4), expected=bias2a) 

1574 assertLookup(detector=2, timespan=Timespan(t1, t5), expected=bias2a) 

1575 assertLookup(detector=2, timespan=Timespan(t1, None), expected=Ambiguous) 

1576 assertLookup(detector=2, timespan=Timespan(t2, t3), expected=bias2a) 

1577 assertLookup(detector=2, timespan=Timespan(t2, t4), expected=bias2a) 

1578 assertLookup(detector=2, timespan=Timespan(t2, t5), expected=bias2a) 

1579 assertLookup(detector=2, timespan=Timespan(t2, None), expected=Ambiguous) 

1580 assertLookup(detector=2, timespan=Timespan(t3, t4), expected=None) 

1581 assertLookup(detector=2, timespan=Timespan(t3, t5), expected=None) 

1582 assertLookup(detector=2, timespan=Timespan(t3, None), expected=bias2b) 

1583 assertLookup(detector=2, timespan=Timespan(t4, t5), expected=None) 

1584 assertLookup(detector=2, timespan=Timespan(t4, None), expected=bias2b) 

1585 assertLookup(detector=2, timespan=Timespan(t5, None), expected=bias2b) 

1586 assertLookup(detector=3, timespan=Timespan(None, t1), expected=None) 

1587 assertLookup(detector=3, timespan=Timespan(None, t2), expected=bias3a) 

1588 assertLookup(detector=3, timespan=Timespan(None, t3), expected=bias3a) 

1589 assertLookup(detector=3, timespan=Timespan(None, t4), expected=bias3a) 

1590 assertLookup(detector=3, timespan=Timespan(None, t5), expected=bias3a) 

1591 assertLookup(detector=3, timespan=Timespan(None, None), expected=Ambiguous) 

1592 assertLookup(detector=3, timespan=Timespan(t1, t2), expected=bias3a) 

1593 assertLookup(detector=3, timespan=Timespan(t1, t3), expected=bias3a) 

1594 assertLookup(detector=3, timespan=Timespan(t1, t4), expected=bias3a) 

1595 assertLookup(detector=3, timespan=Timespan(t1, t5), expected=bias3a) 

1596 assertLookup(detector=3, timespan=Timespan(t1, None), expected=Ambiguous) 

1597 assertLookup(detector=3, timespan=Timespan(t2, t3), expected=bias3a) 

1598 assertLookup(detector=3, timespan=Timespan(t2, t4), expected=bias3a) 

1599 assertLookup(detector=3, timespan=Timespan(t2, t5), expected=bias3a) 

1600 assertLookup(detector=3, timespan=Timespan(t2, None), expected=Ambiguous) 

1601 assertLookup(detector=3, timespan=Timespan(t3, t4), expected=None) 

1602 assertLookup(detector=3, timespan=Timespan(t3, t5), expected=None) 

1603 assertLookup(detector=3, timespan=Timespan(t3, None), expected=bias3b) 

1604 assertLookup(detector=3, timespan=Timespan(t4, t5), expected=None) 

1605 assertLookup(detector=3, timespan=Timespan(t4, None), expected=bias3b) 

1606 assertLookup(detector=3, timespan=Timespan(t5, None), expected=bias3b) 

1607 

1608 # Decertify everything, this time with explicit data IDs, then check 

1609 # that no lookups succeed. 

1610 registry.decertify( 

1611 collection, "bias", Timespan(None, None), 

1612 dataIds=[ 

1613 dict(instrument="Cam1", detector=2), 

1614 dict(instrument="Cam1", detector=3), 

1615 ] 

1616 ) 

1617 for detector in (2, 3): 

1618 for timespan in allTimespans: 

1619 assertLookup(detector=detector, timespan=timespan, expected=None) 

1620 # Certify bias2a and bias3a over (-∞, ∞), check that all lookups return 

1621 # those. 

1622 registry.certify(collection, [bias2a, bias3a], Timespan(None, None),) 

1623 for timespan in allTimespans: 

1624 assertLookup(detector=2, timespan=timespan, expected=bias2a) 

1625 assertLookup(detector=3, timespan=timespan, expected=bias3a) 

1626 # Decertify just bias2 over [t2, t4). 

1627 # This should split a single certification row into two (and leave the 

1628 # other existing row, for bias3a, alone). 

1629 registry.decertify(collection, "bias", Timespan(t2, t4), 

1630 dataIds=[dict(instrument="Cam1", detector=2)]) 

1631 for timespan in allTimespans: 

1632 assertLookup(detector=3, timespan=timespan, expected=bias3a) 

1633 overlapsBefore = timespan.overlaps(Timespan(None, t2)) 

1634 overlapsAfter = timespan.overlaps(Timespan(t4, None)) 

1635 if overlapsBefore and overlapsAfter: 

1636 expected = Ambiguous 

1637 elif overlapsBefore or overlapsAfter: 

1638 expected = bias2a 

1639 else: 

1640 expected = None 

1641 assertLookup(detector=2, timespan=timespan, expected=expected)