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 os 

27 

28import astropy.time 

29import sqlalchemy 

30 

31from ...core import ( 

32 DataCoordinate, 

33 DatasetType, 

34 DimensionGraph, 

35 StorageClass, 

36 ddl, 

37 YamlRepoImportBackend 

38) 

39from .._registry import ( 

40 CollectionType, 

41 ConflictingDefinitionError, 

42 ConsistentDataIds, 

43 OrphanedRecordError, 

44 Registry, 

45) 

46from ..wildcards import DatasetTypeRestriction 

47from ..interfaces import MissingCollectionError 

48 

49 

50class RegistryTests(ABC): 

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

52 generate tests for different configurations. 

53 """ 

54 

55 @classmethod 

56 @abstractmethod 

57 def getDataDir(cls) -> str: 

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

59 """ 

60 raise NotImplementedError() 

61 

62 @abstractmethod 

63 def makeRegistry(self) -> Registry: 

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

65 """ 

66 raise NotImplementedError() 

67 

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

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

70 which should be a YAML import/export file. 

71 """ 

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

73 backend = YamlRepoImportBackend(stream, registry) 

74 backend.register() 

75 backend.load(datastore=None) 

76 

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

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

79 """ 

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

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

82 # of working. 

83 sql = sqlalchemy.sql.select( 

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

85 ).select_from( 

86 getattr(registry._tables, table) 

87 ) 

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

89 

90 def testOpaque(self): 

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

92 `Registry.insertOpaqueData`, `Registry.fetchOpaqueData`, and 

93 `Registry.deleteOpaqueData`. 

94 """ 

95 registry = self.makeRegistry() 

96 table = "opaque_table_for_testing" 

97 registry.registerOpaqueTable( 

98 table, 

99 spec=ddl.TableSpec( 

100 fields=[ 

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

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

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

104 ], 

105 ) 

106 ) 

107 rows = [ 

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

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

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

111 ] 

112 registry.insertOpaqueData(table, *rows) 

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

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

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

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

117 registry.deleteOpaqueData(table, id=3) 

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

119 registry.deleteOpaqueData(table) 

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

121 

122 def testDatasetType(self): 

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

124 `Registry.getDatasetType`. 

125 """ 

126 registry = self.makeRegistry() 

127 # Check valid insert 

128 datasetTypeName = "test" 

129 storageClass = StorageClass("testDatasetType") 

130 registry.storageClasses.registerStorageClass(storageClass) 

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

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

133 inDatasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

134 # Inserting for the first time should return True 

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

136 outDatasetType1 = registry.getDatasetType(datasetTypeName) 

137 self.assertEqual(outDatasetType1, inDatasetType) 

138 

139 # Re-inserting should work 

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

141 # Except when they are not identical 

142 with self.assertRaises(ConflictingDefinitionError): 

143 nonIdenticalDatasetType = DatasetType(datasetTypeName, differentDimensions, storageClass) 

144 registry.registerDatasetType(nonIdenticalDatasetType) 

145 

146 # Template can be None 

147 datasetTypeName = "testNoneTemplate" 

148 storageClass = StorageClass("testDatasetType2") 

149 registry.storageClasses.registerStorageClass(storageClass) 

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

151 inDatasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

152 registry.registerDatasetType(inDatasetType) 

153 outDatasetType2 = registry.getDatasetType(datasetTypeName) 

154 self.assertEqual(outDatasetType2, inDatasetType) 

155 

156 allTypes = set(registry.queryDatasetTypes()) 

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

158 

159 def testDimensions(self): 

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

161 `Registry.syncDimensionData`, and `Registry.expandDataId`. 

162 """ 

163 registry = self.makeRegistry() 

164 dimensionName = "instrument" 

165 dimension = registry.dimensions[dimensionName] 

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

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

168 registry.insertDimensionData(dimensionName, dimensionValue) 

169 # Inserting the same value twice should fail 

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

171 registry.insertDimensionData(dimensionName, dimensionValue) 

172 # expandDataId should retrieve the record we just inserted 

173 self.assertEqual( 

174 registry.expandDataId( 

175 instrument="DummyCam", 

176 graph=dimension.graph 

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

178 dimensionValue 

179 ) 

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

181 with self.assertRaises(LookupError): 

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

183 # abstract_filter doesn't have a table; insert should fail. 

184 with self.assertRaises(TypeError): 

185 registry.insertDimensionData("abstract_filter", {"abstract_filter": "i"}) 

186 dimensionName2 = "physical_filter" 

187 dimension2 = registry.dimensions[dimensionName2] 

188 dimensionValue2 = {"name": "DummyCam_i", "abstract_filter": "i"} 

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

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

191 registry.insertDimensionData(dimensionName2, dimensionValue2) 

192 # Adding required dependency should fix the failure 

193 dimensionValue2["instrument"] = "DummyCam" 

194 registry.insertDimensionData(dimensionName2, dimensionValue2) 

195 # expandDataId should retrieve the record we just inserted. 

196 self.assertEqual( 

197 registry.expandDataId( 

198 instrument="DummyCam", physical_filter="DummyCam_i", 

199 graph=dimension2.graph 

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

201 dimensionValue2 

202 ) 

203 # Use syncDimensionData to insert a new record successfully. 

204 dimensionName3 = "detector" 

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

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

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

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

209 # should be okay. 

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

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

212 # This should fail. 

213 with self.assertRaises(ConflictingDefinitionError): 

214 registry.syncDimensionData( 

215 dimensionName3, 

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

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

218 ) 

219 

220 def testDataIdRelationships(self): 

221 """Test `Registry.relateDataId`. 

222 """ 

223 registry = self.makeRegistry() 

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

225 # Simple cases where the dimension key-value pairs tell us everything. 

226 self.assertEqual( 

227 registry.relateDataIds( 

228 {"instrument": "Cam1"}, 

229 {"instrument": "Cam1"}, 

230 ), 

231 ConsistentDataIds(contains=True, within=True, overlaps=True) 

232 ) 

233 self.assertEqual( 

234 registry.relateDataIds({}, {}), 

235 ConsistentDataIds(contains=True, within=True, overlaps=False) 

236 ) 

237 self.assertEqual( 

238 registry.relateDataIds({"instrument": "Cam1"}, {}), 

239 ConsistentDataIds(contains=True, within=False, overlaps=False) 

240 ) 

241 self.assertEqual( 

242 registry.relateDataIds({}, {"instrument": "Cam1"}), 

243 ConsistentDataIds(contains=False, within=True, overlaps=False) 

244 ) 

245 self.assertEqual( 

246 registry.relateDataIds( 

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

248 {"instrument": "Cam1"}, 

249 ), 

250 ConsistentDataIds(contains=True, within=False, overlaps=True) 

251 ) 

252 self.assertEqual( 

253 registry.relateDataIds( 

254 {"instrument": "Cam1"}, 

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

256 ), 

257 ConsistentDataIds(contains=False, within=True, overlaps=True) 

258 ) 

259 self.assertIsNone( 

260 registry.relateDataIds( 

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

262 {"instrument": "Cam1", "physical_filter": "Cam1-R1"}, 

263 ) 

264 ) 

265 # Trickier cases where we need to expand data IDs, but it's still just 

266 # required and implied dimension relationships. 

267 self.assertEqual( 

268 registry.relateDataIds( 

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

270 {"instrument": "Cam1", "abstract_filter": "g"}, 

271 ), 

272 ConsistentDataIds(contains=True, within=False, overlaps=True) 

273 ) 

274 self.assertEqual( 

275 registry.relateDataIds( 

276 {"instrument": "Cam1", "abstract_filter": "g"}, 

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

278 ), 

279 ConsistentDataIds(contains=False, within=True, overlaps=True) 

280 ) 

281 self.assertEqual( 

282 registry.relateDataIds( 

283 {"instrument": "Cam1"}, 

284 {"htm7": 131073}, 

285 ), 

286 ConsistentDataIds(contains=False, within=False, overlaps=False) 

287 ) 

288 # Trickiest cases involve spatial or temporal overlaps or non-dimension 

289 # elements that relate things (of which visit_definition is our only 

290 # current example). 

291 # 

292 # These two HTM IDs at different levels have a "contains" relationship 

293 # spatially, but there is no overlap of dimension keys. The exact 

294 # result of relateDataIds is unspecified for this case, but it's 

295 # guaranteed to be truthy (see relateDataIds docs.). 

296 self.assertTrue( 

297 registry.relateDataIds({"htm7": 131073}, {"htm9": 2097169}) 

298 ) 

299 # These two HTM IDs at different levels are disjoint spatially, which 

300 # means the data IDs are inconsistent. 

301 self.assertIsNone( 

302 registry.relateDataIds({"htm7": 131073}, {"htm9": 2097391}) 

303 ) 

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

305 registry.insertDimensionData( 

306 "exposure", 

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

308 ) 

309 registry.insertDimensionData( 

310 "exposure", 

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

312 ) 

313 registry.insertDimensionData( 

314 "visit_system", 

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

316 ) 

317 registry.insertDimensionData( 

318 "visit", 

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

320 ) 

321 registry.insertDimensionData( 

322 "visit_definition", 

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

324 ) 

325 self.assertEqual( 

326 registry.relateDataIds( 

327 {"instrument": "Cam1", "visit": 1}, 

328 {"instrument": "Cam1", "exposure": 1}, 

329 ), 

330 ConsistentDataIds(contains=False, within=False, overlaps=True) 

331 ) 

332 self.assertIsNone( 

333 registry.relateDataIds( 

334 {"instrument": "Cam1", "visit": 1}, 

335 {"instrument": "Cam1", "exposure": 2}, 

336 ) 

337 ) 

338 

339 def testDataset(self): 

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

341 and `Registry.removeDatasets`. 

342 """ 

343 registry = self.makeRegistry() 

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

345 run = "test" 

346 registry.registerRun(run) 

347 datasetType = registry.getDatasetType("permabias") 

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

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

350 outRef = registry.getDataset(ref.id) 

351 self.assertIsNotNone(ref.id) 

352 self.assertEqual(ref, outRef) 

353 with self.assertRaises(ConflictingDefinitionError): 

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

355 registry.removeDatasets([ref]) 

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

357 

358 def testComponents(self): 

359 """Tests for `Registry.attachComponent` and other dataset operations 

360 on composite datasets. 

361 """ 

362 registry = self.makeRegistry() 

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

364 run = "test" 

365 registry.registerRun(run) 

366 parentDatasetType = registry.getDatasetType("permabias") 

367 childDatasetType1 = registry.getDatasetType("permabias.image") 

368 childDatasetType2 = registry.getDatasetType("permabias.mask") 

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

370 parent, = registry.insertDatasets(parentDatasetType, dataIds=[dataId], run=run) 

371 children = {"image": registry.insertDatasets(childDatasetType1, dataIds=[dataId], run=run)[0], 

372 "mask": registry.insertDatasets(childDatasetType2, dataIds=[dataId], run=run)[0]} 

373 for name, child in children.items(): 

374 registry.attachComponent(name, parent, child) 

375 self.assertEqual(parent.components, children) 

376 outParent = registry.getDataset(parent.id) 

377 self.assertEqual(outParent.components, children) 

378 # Remove the parent; this should remove all children. 

379 registry.removeDatasets([parent]) 

380 self.assertIsNone(registry.findDataset(parentDatasetType, dataId, collections=[run])) 

381 self.assertIsNone(registry.findDataset(childDatasetType1, dataId, collections=[run])) 

382 self.assertIsNone(registry.findDataset(childDatasetType2, dataId, collections=[run])) 

383 

384 def testFindDataset(self): 

385 """Tests for `Registry.findDataset`. 

386 """ 

387 registry = self.makeRegistry() 

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

389 run = "test" 

390 datasetType = registry.getDatasetType("permabias") 

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

392 registry.registerRun(run) 

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

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

395 self.assertEqual(outputRef, inputRef) 

396 # Check that retrieval with invalid dataId raises 

397 with self.assertRaises(LookupError): 

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

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

400 # Check that different dataIds match to different datasets 

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

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

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

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

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

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

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

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

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

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

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

412 

413 def testCollections(self): 

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

415 """ 

416 registry = self.makeRegistry() 

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

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

419 run1 = "imported_g" 

420 run2 = "imported_r" 

421 datasetType = "permabias" 

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

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

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

425 self.assertIsNotNone(ref1) 

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

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

428 self.assertIsNotNone(ref2) 

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

430 tag1 = "tag1" 

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

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

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

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

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

436 registry.disassociate(tag1, [ref1]) 

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

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

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

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

441 collections = set(registry.queryCollections()) 

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

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

444 # should be a harmless no-op. 

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

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

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

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

449 # dataset type and data ID as ref2. 

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

451 self.assertNotEqual(ref2, ref2b) 

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

453 with self.assertRaises(ConflictingDefinitionError): 

454 registry.associate(tag1, [ref2b]) 

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

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

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

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

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

460 # Should also fail without modifying anything. 

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

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

463 with self.assertRaises(ConflictingDefinitionError): 

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

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

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

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

468 # Register a chained collection that searches: 

469 # 1. 'tag1' 

470 # 2. 'run1', but only for the permaflat dataset 

471 # 3. 'run2' 

472 chain1 = "chain1" 

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

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

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

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

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

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

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

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

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

482 # should fail. 

483 with self.assertRaises(ValueError): 

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

485 # Add the child collections. 

486 registry.setCollectionChain(chain1, [tag1, (run1, "permaflat"), run2]) 

487 self.assertEqual( 

488 list(registry.getCollectionChain(chain1)), 

489 [(tag1, DatasetTypeRestriction.any), 

490 (run1, DatasetTypeRestriction.fromExpression("permaflat")), 

491 (run2, DatasetTypeRestriction.any)] 

492 ) 

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

494 # ref2, because both are in tag1. 

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

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

497 # Now disassociate ref2 from tag1. The search (for permabias) with 

498 # dataId2 in chain1 should then: 

499 # 1. not find it in tag1 

500 # 2. not look in tag2, because it's restricted to permaflat here 

501 # 3. find a different dataset in run2 

502 registry.disassociate(tag1, [ref2]) 

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

504 self.assertNotEqual(ref2b, ref2) 

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

506 # Look in the chain for a permaflat that is in run1; should get the 

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

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

509 self.assertEqual(registry.findDataset("permaflat", dataId3, collections=chain1), 

510 registry.findDataset("permaflat", dataId3, collections=run1),) 

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

512 chain2 = "chain2" 

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

514 registry.setCollectionChain(chain2, [(run2, "permabias"), chain1]) 

515 # Search for permabias with dataId1 should find it via tag1 in chain2, 

516 # recursing, because is not in run1. 

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

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

519 # Search for permabias with dataId2 should find it in run2 (ref2b). 

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

521 # Search for a permaflat that is in run2. That should not be found 

522 # at the front of chain2, because of the restriction to permabias 

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

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

525 ref4 = registry.findDataset("permaflat", dataId4, collections=run2) 

526 self.assertIsNotNone(ref4) 

527 self.assertEqual(ref4, registry.findDataset("permaflat", dataId4, collections=chain2)) 

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

529 # allowed, and is exception-safe. 

530 with self.assertRaises(Exception): 

531 registry.removeCollection(run2) 

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

533 with self.assertRaises(Exception): 

534 registry.removeCollection(chain1) 

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

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

537 registry.removeCollection(chain2) 

538 with self.assertRaises(MissingCollectionError): 

539 registry.getCollectionType(chain2) 

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

541 registry.removeCollection(chain1) 

542 registry.removeCollection(run2) 

543 with self.assertRaises(MissingCollectionError): 

544 registry.getCollectionType(run2) 

545 with self.assertRaises(MissingCollectionError): 

546 registry.getCollectionType(chain1) 

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

548 # collections. 

549 registry.removeCollection(tag1) 

550 with self.assertRaises(MissingCollectionError): 

551 registry.getCollectionType(tag1) 

552 

553 def testDatasetLocations(self): 

554 """Tests for `Registry.insertDatasetLocations`, 

555 `Registry.getDatasetLocations`, and `Registry.removeDatasetLocation`. 

556 """ 

557 registry = self.makeRegistry() 

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

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

560 run = "imported_g" 

561 ref = registry.findDataset("permabias", dataId={"instrument": "Cam1", "detector": 1}, collections=run) 

562 ref2 = registry.findDataset("permaflat", 

563 dataId={"instrument": "Cam1", "detector": 3, "physical_filter": "Cam1-G"}, 

564 collections=run) 

565 datastoreName = "dummystore" 

566 datastoreName2 = "dummystore2" 

567 # Test adding information about a new dataset 

568 registry.insertDatasetLocations(datastoreName, [ref]) 

569 addresses = registry.getDatasetLocations(ref) 

570 self.assertIn(datastoreName, addresses) 

571 self.assertEqual(len(addresses), 1) 

572 registry.insertDatasetLocations(datastoreName2, [ref, ref2]) 

573 addresses = registry.getDatasetLocations(ref) 

574 self.assertEqual(len(addresses), 2) 

575 self.assertIn(datastoreName, addresses) 

576 self.assertIn(datastoreName2, addresses) 

577 registry.removeDatasetLocation(datastoreName, [ref]) 

578 addresses = registry.getDatasetLocations(ref) 

579 self.assertEqual(len(addresses), 1) 

580 self.assertNotIn(datastoreName, addresses) 

581 self.assertIn(datastoreName2, addresses) 

582 with self.assertRaises(OrphanedRecordError): 

583 registry.removeDatasets([ref]) 

584 registry.removeDatasetLocation(datastoreName2, [ref]) 

585 addresses = registry.getDatasetLocations(ref) 

586 self.assertEqual(len(addresses), 0) 

587 self.assertNotIn(datastoreName2, addresses) 

588 registry.removeDatasets([ref]) # should not raise 

589 addresses = registry.getDatasetLocations(ref2) 

590 self.assertEqual(len(addresses), 1) 

591 self.assertIn(datastoreName2, addresses) 

592 

593 def testBasicTransaction(self): 

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

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

596 """ 

597 registry = self.makeRegistry() 

598 storageClass = StorageClass("testDatasetType") 

599 registry.storageClasses.registerStorageClass(storageClass) 

600 dimensions = registry.dimensions.extract(("instrument",)) 

601 dataId = {"instrument": "DummyCam"} 

602 datasetTypeA = DatasetType(name="A", 

603 dimensions=dimensions, 

604 storageClass=storageClass) 

605 datasetTypeB = DatasetType(name="B", 

606 dimensions=dimensions, 

607 storageClass=storageClass) 

608 datasetTypeC = DatasetType(name="C", 

609 dimensions=dimensions, 

610 storageClass=storageClass) 

611 run = "test" 

612 registry.registerRun(run) 

613 refId = None 

614 with registry.transaction(): 

615 registry.registerDatasetType(datasetTypeA) 

616 with self.assertRaises(ValueError): 

617 with registry.transaction(): 

618 registry.registerDatasetType(datasetTypeB) 

619 registry.registerDatasetType(datasetTypeC) 

620 registry.insertDimensionData("instrument", {"instrument": "DummyCam"}) 

621 ref, = registry.insertDatasets(datasetTypeA, dataIds=[dataId], run=run) 

622 refId = ref.id 

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

624 # A should exist 

625 self.assertEqual(registry.getDatasetType("A"), datasetTypeA) 

626 # But B and C should both not exist 

627 with self.assertRaises(KeyError): 

628 registry.getDatasetType("B") 

629 with self.assertRaises(KeyError): 

630 registry.getDatasetType("C") 

631 # And neither should the dataset 

632 self.assertIsNotNone(refId) 

633 self.assertIsNone(registry.getDataset(refId)) 

634 # Or the Dimension entries 

635 with self.assertRaises(LookupError): 

636 registry.expandDataId({"instrument": "DummyCam"}) 

637 

638 def testNestedTransaction(self): 

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

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

641 then caught. 

642 """ 

643 registry = self.makeRegistry() 

644 dimension = registry.dimensions["instrument"] 

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

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

647 checkpointReached = False 

648 with registry.transaction(): 

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

650 registry.insertDimensionData(dimension, dataId1) 

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

652 with registry.transaction(): 

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

654 # be committed). 

655 registry.insertDimensionData(dimension, dataId2) 

656 checkpointReached = True 

657 # This should conflict and raise, triggerring a rollback 

658 # of the previous insertion within the same transaction 

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

660 # block. 

661 registry.insertDimensionData(dimension, dataId1) 

662 self.assertTrue(checkpointReached) 

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

664 with self.assertRaises(LookupError): 

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

666 

667 def testInstrumentDimensions(self): 

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

669 skymap.""" 

670 registry = self.makeRegistry() 

671 

672 # need a bunch of dimensions and datasets for test 

673 registry.insertDimensionData( 

674 "instrument", 

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

676 ) 

677 registry.insertDimensionData( 

678 "physical_filter", 

679 dict(instrument="DummyCam", name="dummy_r", abstract_filter="r"), 

680 dict(instrument="DummyCam", name="dummy_i", abstract_filter="i"), 

681 ) 

682 registry.insertDimensionData( 

683 "detector", 

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

685 ) 

686 registry.insertDimensionData( 

687 "visit_system", 

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

689 ) 

690 registry.insertDimensionData( 

691 "visit", 

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

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

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

695 ) 

696 registry.insertDimensionData( 

697 "exposure", 

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

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

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

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

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

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

704 ) 

705 registry.insertDimensionData( 

706 "visit_definition", 

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

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

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

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

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

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

713 ) 

714 # dataset types 

715 run1 = "test1_r" 

716 run2 = "test2_r" 

717 tagged2 = "test2_t" 

718 registry.registerRun(run1) 

719 registry.registerRun(run2) 

720 registry.registerCollection(tagged2) 

721 storageClass = StorageClass("testDataset") 

722 registry.storageClasses.registerStorageClass(storageClass) 

723 rawType = DatasetType(name="RAW", 

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

725 storageClass=storageClass) 

726 registry.registerDatasetType(rawType) 

727 calexpType = DatasetType(name="CALEXP", 

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

729 storageClass=storageClass) 

730 registry.registerDatasetType(calexpType) 

731 

732 # add pre-existing datasets 

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

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

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=run1) 

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

739 # 100 has different datasets in the different collections 

740 # 101 has the same dataset in both collections. 

741 if exposure == 100: 

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

743 if exposure in (100, 101): 

744 registry.associate(tagged2, [ref]) 

745 # Add pre-existing datasets to tagged2. 

746 for exposure in (200, 201): 

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

748 # note that only 3 of 5 detectors have datasets 

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

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

751 registry.associate(tagged2, [ref]) 

752 

753 dimensions = DimensionGraph( 

754 registry.dimensions, 

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

756 ) 

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

758 rows = list(registry.queryDimensions("visit", datasets=rawType, collections=run1, expand=True)) 

759 rowsI = list(registry.queryDimensions(["visit"], datasets=rawType, collections=run1, expand=True)) 

760 self.assertEqual(rows, rowsI) 

761 # with empty expression 

762 rows = list(registry.queryDimensions(dimensions, datasets=rawType, collections=run1, expand=True)) 

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

764 for dataId in rows: 

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

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

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

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

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

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

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

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

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

774 (100, 101, 110, 111)) 

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

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

777 

778 # second collection 

779 rows = list(registry.queryDimensions(dimensions, datasets=rawType, collections=tagged2)) 

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

781 for dataId in rows: 

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

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

784 (100, 101, 200, 201)) 

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

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

787 

788 # with two input datasets 

789 rows = list(registry.queryDimensions(dimensions, datasets=rawType, collections=[run1, tagged2])) 

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

791 for dataId in rows: 

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

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

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

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

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

797 

798 # limit to single visit 

799 rows = list(registry.queryDimensions(dimensions, datasets=rawType, collections=run1, 

800 where="visit = 10")) 

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

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

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

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

805 

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

807 rows = list(registry.queryDimensions(dimensions, datasets=rawType, collections=run1, 

808 where="visit = 10 and detector > 1")) 

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

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

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

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

813 

814 # expression excludes everything 

815 rows = list(registry.queryDimensions(dimensions, datasets=rawType, collections=run1, 

816 where="visit > 1000")) 

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

818 

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

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

821 rows = list(registry.queryDimensions(dimensions, datasets=rawType, collections=run1, 

822 where="physical_filter = 'dummy_r'")) 

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

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

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

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

827 

828 def testSkyMapDimensions(self): 

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

830 registry = self.makeRegistry() 

831 

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

833 # "abstract_filter" in the test so also have to add physical_filter 

834 # dimensions 

835 registry.insertDimensionData( 

836 "instrument", 

837 dict(instrument="DummyCam") 

838 ) 

839 registry.insertDimensionData( 

840 "physical_filter", 

841 dict(instrument="DummyCam", name="dummy_r", abstract_filter="r"), 

842 dict(instrument="DummyCam", name="dummy_i", abstract_filter="i"), 

843 ) 

844 registry.insertDimensionData( 

845 "skymap", 

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

847 ) 

848 for tract in range(10): 

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

850 registry.insertDimensionData( 

851 "patch", 

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

853 for patch in range(10)] 

854 ) 

855 

856 # dataset types 

857 run = "test" 

858 registry.registerRun(run) 

859 storageClass = StorageClass("testDataset") 

860 registry.storageClasses.registerStorageClass(storageClass) 

861 calexpType = DatasetType(name="deepCoadd_calexp", 

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

863 "abstract_filter")), 

864 storageClass=storageClass) 

865 registry.registerDatasetType(calexpType) 

866 mergeType = DatasetType(name="deepCoadd_mergeDet", 

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

868 storageClass=storageClass) 

869 registry.registerDatasetType(mergeType) 

870 measType = DatasetType(name="deepCoadd_meas", 

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

872 "abstract_filter")), 

873 storageClass=storageClass) 

874 registry.registerDatasetType(measType) 

875 

876 dimensions = DimensionGraph( 

877 registry.dimensions, 

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

879 | measType.dimensions.required) 

880 ) 

881 

882 # add pre-existing datasets 

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

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

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

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

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

888 dataId = dict(skymap="DummyMap", tract=tract, patch=patch, abstract_filter=aFilter) 

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

890 

891 # with empty expression 

892 rows = list(registry.queryDimensions(dimensions, 

893 datasets=[calexpType, mergeType], collections=run)) 

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

895 for dataId in rows: 

896 self.assertCountEqual(dataId.keys(), ("skymap", "tract", "patch", "abstract_filter")) 

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

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

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

900 

901 # limit to 2 tracts and 2 patches 

902 rows = list(registry.queryDimensions(dimensions, 

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

904 where="tract IN (1, 5) AND patch IN (2, 7)")) 

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

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

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

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

909 

910 # limit to single filter 

911 rows = list(registry.queryDimensions(dimensions, 

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

913 where="abstract_filter = 'i'")) 

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

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

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

917 self.assertCountEqual(set(dataId["abstract_filter"] for dataId in rows), ("i",)) 

918 

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

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

921 rows = list(registry.queryDimensions(dimensions, 

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

923 where="skymap = 'Mars'")) 

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

925 

926 def testSpatialMatch(self): 

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

928 

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

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

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

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

933 that generates query. 

934 """ 

935 registry = self.makeRegistry() 

936 

937 # dataset types 

938 collection = "test" 

939 registry.registerRun(name=collection) 

940 storageClass = StorageClass("testDataset") 

941 registry.storageClasses.registerStorageClass(storageClass) 

942 

943 calexpType = DatasetType(name="CALEXP", 

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

945 storageClass=storageClass) 

946 registry.registerDatasetType(calexpType) 

947 

948 coaddType = DatasetType(name="deepCoadd_calexp", 

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

950 "abstract_filter")), 

951 storageClass=storageClass) 

952 registry.registerDatasetType(coaddType) 

953 

954 dimensions = DimensionGraph( 

955 registry.dimensions, 

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

957 ) 

958 

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

960 rows = list(registry.queryDimensions(dimensions, datasets=calexpType, collections=collection)) 

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

962 

963 def testCalibrationLabelIndirection(self): 

964 """Test that we can look up datasets with calibration_label dimensions 

965 from a data ID with exposure dimensions. 

966 """ 

967 

968 def _dt(iso_string): 

969 return astropy.time.Time(iso_string, format="iso", scale="utc") 

970 

971 registry = self.makeRegistry() 

972 

973 flat = DatasetType( 

974 "flat", 

975 registry.dimensions.extract( 

976 ["instrument", "detector", "physical_filter", "calibration_label"] 

977 ), 

978 "ImageU" 

979 ) 

980 registry.registerDatasetType(flat) 

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

982 registry.insertDimensionData( 

983 "physical_filter", 

984 dict(instrument="DummyCam", name="dummy_i", abstract_filter="i"), 

985 ) 

986 registry.insertDimensionData( 

987 "detector", 

988 *[dict(instrument="DummyCam", id=i, full_name=str(i)) for i in (1, 2, 3, 4, 5)] 

989 ) 

990 registry.insertDimensionData( 

991 "exposure", 

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

993 datetime_begin=_dt("2005-12-15 02:00:00"), datetime_end=_dt("2005-12-15 03:00:00")), 

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

995 datetime_begin=_dt("2005-12-16 02:00:00"), datetime_end=_dt("2005-12-16 03:00:00")), 

996 ) 

997 registry.insertDimensionData( 

998 "calibration_label", 

999 dict(instrument="DummyCam", name="first_night", 

1000 datetime_begin=_dt("2005-12-15 01:00:00"), datetime_end=_dt("2005-12-15 04:00:00")), 

1001 dict(instrument="DummyCam", name="second_night", 

1002 datetime_begin=_dt("2005-12-16 01:00:00"), datetime_end=_dt("2005-12-16 04:00:00")), 

1003 dict(instrument="DummyCam", name="both_nights", 

1004 datetime_begin=_dt("2005-12-15 01:00:00"), datetime_end=_dt("2005-12-16 04:00:00")), 

1005 ) 

1006 # Different flats for different nights for detectors 1-3 in first 

1007 # collection. 

1008 run1 = "calibs1" 

1009 registry.registerRun(run1) 

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

1011 registry.insertDatasets(flat, [dict(instrument="DummyCam", calibration_label="first_night", 

1012 physical_filter="dummy_i", detector=detector)], 

1013 run=run1) 

1014 registry.insertDatasets(flat, [dict(instrument="DummyCam", calibration_label="second_night", 

1015 physical_filter="dummy_i", detector=detector)], 

1016 run=run1) 

1017 # The same flat for both nights for detectors 3-5 (so detector 3 has 

1018 # multiple valid flats) in second collection. 

1019 run2 = "calib2" 

1020 registry.registerRun(run2) 

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

1022 registry.insertDatasets(flat, [dict(instrument="DummyCam", calibration_label="both_nights", 

1023 physical_filter="dummy_i", detector=detector)], 

1024 run=run2) 

1025 # Perform queries for individual exposure+detector combinations, which 

1026 # should always return exactly one flat. 

1027 for exposure in (100, 101): 

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

1029 with self.subTest(exposure=exposure, detector=detector): 

1030 rows = list(registry.queryDatasets("flat", collections=[run1], 

1031 instrument="DummyCam", 

1032 exposure=exposure, 

1033 detector=detector)) 

1034 self.assertEqual(len(rows), 1) 

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

1036 with self.subTest(exposure=exposure, detector=detector): 

1037 rows = registry.queryDatasets("flat", collections=[run2], 

1038 instrument="DummyCam", 

1039 exposure=exposure, 

1040 detector=detector) 

1041 self.assertEqual(len(list(rows)), 1) 

1042 for detector in (1, 2, 4, 5): 

1043 with self.subTest(exposure=exposure, detector=detector): 

1044 rows = registry.queryDatasets("flat", collections=[run1, run2], 

1045 instrument="DummyCam", 

1046 exposure=exposure, 

1047 detector=detector) 

1048 self.assertEqual(len(list(rows)), 1) 

1049 for detector in (3,): 

1050 with self.subTest(exposure=exposure, detector=detector): 

1051 rows = registry.queryDatasets("flat", collections=[run1, run2], 

1052 instrument="DummyCam", 

1053 exposure=exposure, 

1054 detector=detector) 

1055 self.assertEqual(len(list(rows)), 2) 

1056 

1057 def testAbstractFilterQuery(self): 

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

1059 abstract_filters. This is tricky because abstract_filter is 

1060 backed by a query against physical_filter. 

1061 """ 

1062 registry = self.makeRegistry() 

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

1064 registry.insertDimensionData( 

1065 "physical_filter", 

1066 dict(instrument="DummyCam", name="dummy_i", abstract_filter="i"), 

1067 dict(instrument="DummyCam", name="dummy_i2", abstract_filter="i"), 

1068 dict(instrument="DummyCam", name="dummy_r", abstract_filter="r"), 

1069 ) 

1070 rows = list(registry.queryDimensions(["abstract_filter"])) 

1071 self.assertCountEqual( 

1072 rows, 

1073 [DataCoordinate.standardize(abstract_filter="i", universe=registry.dimensions), 

1074 DataCoordinate.standardize(abstract_filter="r", universe=registry.dimensions)] 

1075 )