Coverage for tests / test_datastore.py: 13%

1211 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-01-13 01:20 -0800

1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This software is dual licensed under the GNU General Public License and also 

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

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

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

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

14# 

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

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

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

18# (at your option) any later version. 

19# 

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

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

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

23# GNU General Public License for more details. 

24# 

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

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

27 

28from __future__ import annotations 

29 

30import contextlib 

31import os 

32import pickle 

33import shutil 

34import tempfile 

35import time 

36import unittest 

37import unittest.mock 

38import uuid 

39from collections.abc import Callable, Iterator 

40from typing import Any, cast 

41 

42import yaml 

43 

44import lsst.utils.tests 

45from lsst.daf.butler import ( 

46 Config, 

47 DataCoordinate, 

48 DatasetIdGenEnum, 

49 DatasetRef, 

50 DatasetType, 

51 DatasetTypeNotSupportedError, 

52 Datastore, 

53 DimensionUniverse, 

54 FileDataset, 

55 StorageClass, 

56 StorageClassFactory, 

57) 

58from lsst.daf.butler.datastore import DatasetRefURIs, DatastoreConfig, DatastoreValidationError, NullDatastore 

59from lsst.daf.butler.datastore.cache_manager import ( 

60 DatastoreCacheManager, 

61 DatastoreCacheManagerConfig, 

62 DatastoreDisabledCacheManager, 

63) 

64from lsst.daf.butler.datastore.record_data import DatastoreRecordData, SerializedDatastoreRecordData 

65from lsst.daf.butler.datastore.stored_file_info import StoredFileInfo, make_datastore_path_relative 

66from lsst.daf.butler.formatters.yaml import YamlFormatter 

67from lsst.daf.butler.tests import ( 

68 BadNoWriteFormatter, 

69 BadWriteFormatter, 

70 DatasetTestHelper, 

71 DatastoreTestHelper, 

72 DummyRegistry, 

73 MetricsExample, 

74 MetricsExampleDataclass, 

75 MetricsExampleModel, 

76) 

77from lsst.daf.butler.tests.dict_convertible_model import DictConvertibleModel 

78from lsst.daf.butler.tests.utils import TestCaseMixin 

79from lsst.resources import ResourcePath 

80from lsst.utils import doImport 

81from lsst.utils.introspection import get_full_type_name 

82 

83TESTDIR = os.path.dirname(__file__) 

84 

85 

86def makeExampleMetrics(use_none: bool = False) -> MetricsExample: 

87 """Make example dataset that can be stored in butler.""" 

88 if use_none: 

89 array = None 

90 else: 

91 array = [563, 234, 456.7, 105, 2054, -1045] 

92 return MetricsExample( 

93 {"AM1": 5.2, "AM2": 30.6}, 

94 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}}, 

95 array, 

96 ) 

97 

98 

99class TransactionTestError(Exception): 

100 """Specific error for transactions, to prevent misdiagnosing 

101 that might otherwise occur when a standard exception is used. 

102 """ 

103 

104 pass 

105 

106 

107class DatastoreTestsBase(DatasetTestHelper, DatastoreTestHelper, TestCaseMixin): 

108 """Support routines for datastore testing""" 

109 

110 root: str | None = None 

111 universe: DimensionUniverse 

112 storageClassFactory: StorageClassFactory 

113 

114 @classmethod 

115 def setUpClass(cls) -> None: 

116 # Storage Classes are fixed for all datastores in these tests 

117 scConfigFile = os.path.join(TESTDIR, "config/basic/storageClasses.yaml") 

118 cls.storageClassFactory = StorageClassFactory() 

119 cls.storageClassFactory.addFromConfig(scConfigFile) 

120 

121 # Read the Datastore config so we can get the class 

122 # information (since we should not assume the constructor 

123 # name here, but rely on the configuration file itself) 

124 datastoreConfig = DatastoreConfig(cls.configFile) 

125 cls.datastoreType = cast(type[Datastore], doImport(datastoreConfig["cls"])) 

126 cls.universe = DimensionUniverse() 

127 

128 def setUp(self) -> None: 

129 self.setUpDatastoreTests(DummyRegistry, DatastoreConfig) 

130 

131 def tearDown(self) -> None: 

132 if self.root is not None and os.path.exists(self.root): 

133 shutil.rmtree(self.root, ignore_errors=True) 

134 

135 

136class DatastoreTests(DatastoreTestsBase): 

137 """Some basic tests of a simple datastore.""" 

138 

139 hasUnsupportedPut = True 

140 rootKeys: tuple[str, ...] | None = None 

141 isEphemeral: bool = False 

142 validationCanFail: bool = False 

143 

144 def testConfigRoot(self) -> None: 

145 full = DatastoreConfig(self.configFile) 

146 config = DatastoreConfig(self.configFile, mergeDefaults=False) 

147 newroot = "/random/location" 

148 self.datastoreType.setConfigRoot(newroot, config, full) 

149 if self.rootKeys: 

150 for k in self.rootKeys: 

151 self.assertIn(newroot, config[k]) 

152 

153 def testConstructor(self) -> None: 

154 datastore = self.makeDatastore() 

155 self.assertIsNotNone(datastore) 

156 self.assertIs(datastore.isEphemeral, self.isEphemeral) 

157 

158 def testConfigurationValidation(self) -> None: 

159 datastore = self.makeDatastore() 

160 sc = self.storageClassFactory.getStorageClass("ThingOne") 

161 datastore.validateConfiguration([sc]) 

162 

163 sc2 = self.storageClassFactory.getStorageClass("ThingTwo") 

164 if self.validationCanFail: 

165 with self.assertRaises(DatastoreValidationError): 

166 datastore.validateConfiguration([sc2], logFailures=True) 

167 

168 dimensions = self.universe.conform(("visit", "physical_filter")) 

169 dataId = { 

170 "instrument": "dummy", 

171 "visit": 52, 

172 "physical_filter": "V", 

173 "band": "v", 

174 "day_obs": 20250101, 

175 } 

176 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

177 datastore.validateConfiguration([ref]) 

178 

179 def testParameterValidation(self) -> None: 

180 """Check that parameters are validated""" 

181 sc = self.storageClassFactory.getStorageClass("ThingOne") 

182 dimensions = self.universe.conform(("visit", "physical_filter")) 

183 dataId = { 

184 "instrument": "dummy", 

185 "visit": 52, 

186 "physical_filter": "V", 

187 "band": "v", 

188 "day_obs": 20250101, 

189 } 

190 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

191 datastore = self.makeDatastore() 

192 data = {1: 2, 3: 4} 

193 datastore.put(data, ref) 

194 newdata = datastore.get(ref) 

195 self.assertEqual(data, newdata) 

196 with self.assertRaises(KeyError): 

197 newdata = datastore.get(ref, parameters={"missing": 5}) 

198 

199 def testBasicPutGet(self) -> None: 

200 metrics = makeExampleMetrics() 

201 datastore = self.makeDatastore() 

202 

203 # Create multiple storage classes for testing different formulations 

204 storageClasses = [ 

205 self.storageClassFactory.getStorageClass(sc) 

206 for sc in ("StructuredData", "StructuredDataJson", "StructuredDataPickle") 

207 ] 

208 

209 dimensions = self.universe.conform(("visit", "physical_filter")) 

210 dataId = { 

211 "instrument": "dummy", 

212 "visit": 52, 

213 "physical_filter": "V", 

214 "band": "v", 

215 "day_obs": 20250101, 

216 } 

217 dataId2 = { 

218 "instrument": "dummy", 

219 "visit": 53, 

220 "physical_filter": "V", 

221 "band": "v", 

222 "day_obs": 20250101, 

223 } 

224 

225 for sc in storageClasses: 

226 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

227 ref2 = self.makeDatasetRef("metric", dimensions, sc, dataId2) 

228 

229 # Make sure that using getManyURIs without predicting before the 

230 # dataset has been put raises. 

231 with self.assertRaises(FileNotFoundError): 

232 datastore.getManyURIs([ref], predict=False) 

233 

234 # Make sure that using getManyURIs with predicting before the 

235 # dataset has been put predicts the URI. 

236 uris = datastore.getManyURIs([ref, ref2], predict=True) 

237 self.assertIn("52", uris[ref].primaryURI.geturl()) 

238 self.assertIn("#predicted", uris[ref].primaryURI.geturl()) 

239 self.assertIn("53", uris[ref2].primaryURI.geturl()) 

240 self.assertIn("#predicted", uris[ref2].primaryURI.geturl()) 

241 

242 datastore.put(metrics, ref) 

243 

244 # Does it exist? 

245 self.assertTrue(datastore.exists(ref)) 

246 self.assertTrue(datastore.knows(ref)) 

247 multi = datastore.knows_these([ref]) 

248 self.assertTrue(multi[ref]) 

249 multi = datastore.mexists([ref, ref2]) 

250 self.assertTrue(multi[ref]) 

251 self.assertFalse(multi[ref2]) 

252 

253 # Get 

254 metricsOut = datastore.get(ref, parameters=None) 

255 self.assertEqual(metrics, metricsOut) 

256 

257 uri = datastore.getURI(ref) 

258 self.assertEqual(uri.scheme, self.uriScheme) 

259 

260 uris = datastore.getManyURIs([ref]) 

261 self.assertEqual(len(uris), 1) 

262 ref, uri = uris.popitem() 

263 self.assertTrue(uri.primaryURI.exists()) 

264 self.assertFalse(uri.componentURIs) 

265 

266 # Get a component -- we need to construct new refs for them 

267 # with derived storage classes but with parent ID 

268 for comp in ("data", "output"): 

269 compRef = ref.makeComponentRef(comp) 

270 output = datastore.get(compRef) 

271 self.assertEqual(output, getattr(metricsOut, comp)) 

272 

273 uri = datastore.getURI(compRef) 

274 self.assertEqual(uri.scheme, self.uriScheme) 

275 

276 uris = datastore.getManyURIs([compRef]) 

277 self.assertEqual(len(uris), 1) 

278 

279 storageClass = sc 

280 

281 # Check that we can put a metric with None in a component and 

282 # get it back as None 

283 metricsNone = makeExampleMetrics(use_none=True) 

284 dataIdNone = { 

285 "instrument": "dummy", 

286 "visit": 54, 

287 "physical_filter": "V", 

288 "band": "v", 

289 "day_obs": 20250101, 

290 } 

291 refNone = self.makeDatasetRef("metric", dimensions, sc, dataIdNone) 

292 datastore.put(metricsNone, refNone) 

293 

294 comp = "data" 

295 for comp in ("data", "output"): 

296 compRef = refNone.makeComponentRef(comp) 

297 output = datastore.get(compRef) 

298 self.assertEqual(output, getattr(metricsNone, comp)) 

299 

300 # Check that a put fails if the dataset type is not supported 

301 if self.hasUnsupportedPut: 

302 sc = StorageClass("UnsupportedSC", pytype=type(metrics)) 

303 ref = self.makeDatasetRef("unsupportedType", dimensions, sc, dataId) 

304 with self.assertRaises(DatasetTypeNotSupportedError): 

305 datastore.put(metrics, ref) 

306 

307 # These should raise 

308 ref = self.makeDatasetRef("metrics", dimensions, storageClass, dataId) 

309 with self.assertRaises(FileNotFoundError): 

310 # non-existing file 

311 datastore.get(ref) 

312 

313 # Get a URI from it 

314 uri = datastore.getURI(ref, predict=True) 

315 self.assertEqual(uri.scheme, self.uriScheme) 

316 

317 with self.assertRaises(FileNotFoundError): 

318 datastore.getURI(ref) 

319 

320 def testTrustGetRequest(self) -> None: 

321 """Check that we can get datasets that registry knows nothing about.""" 

322 datastore = self.makeDatastore() 

323 

324 # Skip test if the attribute is not defined 

325 if not hasattr(datastore, "trustGetRequest"): 

326 return 

327 

328 metrics = makeExampleMetrics() 

329 

330 i = 0 

331 for sc_name in ("StructuredDataNoComponents", "StructuredData", "StructuredComposite"): 

332 i += 1 

333 datasetTypeName = f"test_metric{i}" # Different dataset type name each time. 

334 

335 if sc_name == "StructuredComposite": 

336 disassembled = True 

337 else: 

338 disassembled = False 

339 

340 # Start datastore in default configuration of using registry 

341 datastore.trustGetRequest = False 

342 

343 # Create multiple storage classes for testing with or without 

344 # disassembly 

345 sc = self.storageClassFactory.getStorageClass(sc_name) 

346 dimensions = self.universe.conform(("visit", "physical_filter")) 

347 

348 dataId = { 

349 "instrument": "dummy", 

350 "visit": 52 + i, 

351 "physical_filter": "V", 

352 "band": "v", 

353 "day_obs": 20250101, 

354 } 

355 

356 ref = self.makeDatasetRef(datasetTypeName, dimensions, sc, dataId) 

357 datastore.put(metrics, ref) 

358 

359 # Does it exist? 

360 self.assertTrue(datastore.exists(ref)) 

361 self.assertTrue(datastore.knows(ref)) 

362 multi = datastore.knows_these([ref]) 

363 self.assertTrue(multi[ref]) 

364 multi = datastore.mexists([ref]) 

365 self.assertTrue(multi[ref]) 

366 

367 # Get 

368 metricsOut = datastore.get(ref) 

369 self.assertEqual(metrics, metricsOut) 

370 

371 # Get the URI(s) 

372 primaryURI, componentURIs = datastore.getURIs(ref) 

373 if disassembled: 

374 self.assertIsNone(primaryURI) 

375 self.assertEqual(len(componentURIs), 3) 

376 else: 

377 self.assertIn(datasetTypeName, primaryURI.path) 

378 self.assertFalse(componentURIs) 

379 

380 # Delete registry entry so now we are trusting 

381 datastore.removeStoredItemInfo(ref) 

382 

383 # Now stop trusting and check that things break 

384 datastore.trustGetRequest = False 

385 

386 # Does it exist? 

387 self.assertFalse(datastore.exists(ref)) 

388 self.assertFalse(datastore.knows(ref)) 

389 multi = datastore.knows_these([ref]) 

390 self.assertFalse(multi[ref]) 

391 multi = datastore.mexists([ref]) 

392 self.assertFalse(multi[ref]) 

393 

394 with self.assertRaises(FileNotFoundError): 

395 datastore.get(ref) 

396 

397 if sc_name != "StructuredDataNoComponents": 

398 with self.assertRaises(FileNotFoundError): 

399 datastore.get(ref.makeComponentRef("data")) 

400 

401 # URI should fail unless we ask for prediction 

402 with self.assertRaises(FileNotFoundError): 

403 datastore.getURIs(ref) 

404 

405 predicted_primary, predicted_disassembled = datastore.getURIs(ref, predict=True) 

406 if disassembled: 

407 self.assertIsNone(predicted_primary) 

408 self.assertEqual(len(predicted_disassembled), 3) 

409 for uri in predicted_disassembled.values(): 

410 self.assertEqual(uri.fragment, "predicted") 

411 self.assertIn(datasetTypeName, uri.path) 

412 else: 

413 self.assertIn(datasetTypeName, predicted_primary.path) 

414 self.assertFalse(predicted_disassembled) 

415 self.assertEqual(predicted_primary.fragment, "predicted") 

416 

417 # Now enable registry-free trusting mode 

418 datastore.trustGetRequest = True 

419 

420 # Try again to get it 

421 metricsOut = datastore.get(ref) 

422 self.assertEqual(metricsOut, metrics) 

423 

424 # Does it exist? 

425 self.assertTrue(datastore.exists(ref)) 

426 

427 # Get a component 

428 if sc_name != "StructuredDataNoComponents": 

429 comp = "data" 

430 compRef = ref.makeComponentRef(comp) 

431 output = datastore.get(compRef) 

432 self.assertEqual(output, getattr(metrics, comp)) 

433 

434 # Get the URI -- if we trust this should work even without 

435 # enabling prediction. 

436 primaryURI2, componentURIs2 = datastore.getURIs(ref) 

437 self.assertEqual(primaryURI2, primaryURI) 

438 self.assertEqual(componentURIs2, componentURIs) 

439 

440 # Check for compatible storage class. 

441 if sc_name in ("StructuredDataNoComponents", "StructuredData"): 

442 # Make new dataset ref with compatible storage class. 

443 ref_comp = ref.overrideStorageClass("StructuredDataDictJson") 

444 

445 # Without `set_retrieve_dataset_type_method` it will fail to 

446 # find correct file. 

447 self.assertFalse(datastore.exists(ref_comp)) 

448 with self.assertRaises(FileNotFoundError): 

449 datastore.get(ref_comp) 

450 with self.assertRaises(FileNotFoundError): 

451 datastore.get(ref, storageClass="StructuredDataDictJson") 

452 

453 # Need a special method to generate stored dataset type. 

454 def _stored_dataset_type(name: str, ref: DatasetRef = ref) -> DatasetType: 

455 if name == ref.datasetType.name: 

456 return ref.datasetType 

457 raise ValueError(f"Unexpected dataset type name {ref.datasetType.name}") 

458 

459 datastore.set_retrieve_dataset_type_method(_stored_dataset_type) 

460 

461 # Storage class override with original dataset ref. 

462 metrics_as_dict = datastore.get(ref, storageClass="StructuredDataDictJson") 

463 self.assertIsInstance(metrics_as_dict, dict) 

464 

465 # get() should return a dict now. 

466 metrics_as_dict = datastore.get(ref_comp) 

467 self.assertIsInstance(metrics_as_dict, dict) 

468 

469 # exists() should work as well. 

470 self.assertTrue(datastore.exists(ref_comp)) 

471 

472 datastore.set_retrieve_dataset_type_method(None) 

473 

474 def testDisassembly(self) -> None: 

475 """Test disassembly within datastore.""" 

476 metrics = makeExampleMetrics() 

477 if self.isEphemeral: 

478 # in-memory datastore does not disassemble 

479 return 

480 

481 # Create multiple storage classes for testing different formulations 

482 # of composites. One of these will not disassemble to provide 

483 # a reference. 

484 storageClasses = [ 

485 self.storageClassFactory.getStorageClass(sc) 

486 for sc in ( 

487 "StructuredComposite", 

488 "StructuredCompositeTestA", 

489 "StructuredCompositeTestB", 

490 "StructuredCompositeReadComp", 

491 "StructuredData", # No disassembly 

492 "StructuredCompositeReadCompNoDisassembly", 

493 ) 

494 ] 

495 

496 # Create the test datastore 

497 datastore = self.makeDatastore() 

498 

499 # Dummy dataId 

500 dimensions = self.universe.conform(("visit", "physical_filter")) 

501 dataId = {"instrument": "dummy", "visit": 428, "physical_filter": "R"} 

502 

503 for i, sc in enumerate(storageClasses): 

504 with self.subTest(storageClass=sc.name): 

505 # Create a different dataset type each time round 

506 # so that a test failure in this subtest does not trigger 

507 # a cascade of tests because of file clashes 

508 ref = self.makeDatasetRef(f"metric_comp_{i}", dimensions, sc, dataId) 

509 

510 disassembled = sc.name not in {"StructuredData", "StructuredCompositeReadCompNoDisassembly"} 

511 

512 datastore.put(metrics, ref) 

513 

514 baseURI, compURIs = datastore.getURIs(ref) 

515 if disassembled: 

516 self.assertIsNone(baseURI) 

517 self.assertEqual(set(compURIs), {"data", "output", "summary"}) 

518 else: 

519 self.assertIsNotNone(baseURI) 

520 self.assertEqual(compURIs, {}) 

521 

522 metrics_get = datastore.get(ref) 

523 self.assertEqual(metrics_get, metrics) 

524 

525 # Retrieve the composite with read parameter 

526 stop = 4 

527 metrics_get = datastore.get(ref, parameters={"slice": slice(stop)}) 

528 self.assertEqual(metrics_get.summary, metrics.summary) 

529 self.assertEqual(metrics_get.output, metrics.output) 

530 self.assertEqual(metrics_get.data, metrics.data[:stop]) 

531 

532 # Retrieve a component 

533 data = datastore.get(ref.makeComponentRef("data")) 

534 self.assertEqual(data, metrics.data) 

535 

536 # On supported storage classes attempt to access a read 

537 # only component 

538 if "ReadComp" in sc.name: 

539 cRef = ref.makeComponentRef("counter") 

540 counter = datastore.get(cRef) 

541 self.assertEqual(counter, len(metrics.data)) 

542 

543 counter = datastore.get(cRef, parameters={"slice": slice(stop)}) 

544 self.assertEqual(counter, stop) 

545 

546 datastore.remove(ref) 

547 

548 def prepDeleteTest(self, n_refs: int = 1) -> tuple[Datastore, tuple[DatasetRef, ...]]: 

549 metrics = makeExampleMetrics() 

550 datastore = self.makeDatastore() 

551 # Put 

552 dimensions = self.universe.conform(("visit", "physical_filter")) 

553 sc = self.storageClassFactory.getStorageClass("StructuredData") 

554 refs = [] 

555 for i in range(n_refs): 

556 dataId = { 

557 "instrument": "dummy", 

558 "visit": 638 + i, 

559 "physical_filter": "U", 

560 "band": "u", 

561 "day_obs": 20250101, 

562 } 

563 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

564 datastore.put(metrics, ref) 

565 

566 # Does it exist? 

567 self.assertTrue(datastore.exists(ref)) 

568 

569 # Get 

570 metricsOut = datastore.get(ref) 

571 self.assertEqual(metrics, metricsOut) 

572 refs.append(ref) 

573 

574 return datastore, *refs 

575 

576 def testRemove(self) -> None: 

577 datastore, ref = self.prepDeleteTest() 

578 

579 # Remove 

580 datastore.remove(ref) 

581 

582 # Does it exist? 

583 self.assertFalse(datastore.exists(ref)) 

584 

585 # Do we now get a predicted URI? 

586 uri = datastore.getURI(ref, predict=True) 

587 self.assertEqual(uri.fragment, "predicted") 

588 

589 # Get should now fail 

590 with self.assertRaises(FileNotFoundError): 

591 datastore.get(ref) 

592 # Can only delete once 

593 with self.assertRaises(FileNotFoundError): 

594 datastore.remove(ref) 

595 

596 def testForget(self) -> None: 

597 datastore, ref = self.prepDeleteTest() 

598 

599 # Remove 

600 datastore.forget([ref]) 

601 

602 # Does it exist (as far as we know)? 

603 self.assertFalse(datastore.exists(ref)) 

604 

605 # Do we now get a predicted URI? 

606 uri = datastore.getURI(ref, predict=True) 

607 self.assertEqual(uri.fragment, "predicted") 

608 

609 # Get should now fail 

610 with self.assertRaises(FileNotFoundError): 

611 datastore.get(ref) 

612 

613 # Forgetting again is a silent no-op 

614 datastore.forget([ref]) 

615 

616 # Predicted URI should still point to the file. 

617 self.assertTrue(uri.exists()) 

618 

619 def testTransfer(self) -> None: 

620 metrics = makeExampleMetrics() 

621 

622 dimensions = self.universe.conform(("visit", "physical_filter")) 

623 dataId = { 

624 "instrument": "dummy", 

625 "visit": 2048, 

626 "physical_filter": "Uprime", 

627 "band": "u", 

628 "day_obs": 20250101, 

629 } 

630 

631 sc = self.storageClassFactory.getStorageClass("StructuredData") 

632 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

633 

634 inputDatastore = self.makeDatastore("test_input_datastore") 

635 outputDatastore = self.makeDatastore("test_output_datastore") 

636 

637 inputDatastore.put(metrics, ref) 

638 outputDatastore.transfer(inputDatastore, ref) 

639 

640 metricsOut = outputDatastore.get(ref) 

641 self.assertEqual(metrics, metricsOut) 

642 

643 def testBasicTransaction(self) -> None: 

644 datastore = self.makeDatastore() 

645 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

646 dimensions = self.universe.conform(("visit", "physical_filter")) 

647 nDatasets = 6 

648 dataIds = [ 

649 {"instrument": "dummy", "visit": i, "physical_filter": "V", "band": "v", "day_obs": 20250101} 

650 for i in range(nDatasets) 

651 ] 

652 data = [ 

653 ( 

654 self.makeDatasetRef("metric", dimensions, storageClass, dataId), 

655 makeExampleMetrics(), 

656 ) 

657 for dataId in dataIds 

658 ] 

659 succeed = data[: nDatasets // 2] 

660 fail = data[nDatasets // 2 :] 

661 # All datasets added in this transaction should continue to exist 

662 with datastore.transaction(): 

663 for ref, metrics in succeed: 

664 datastore.put(metrics, ref) 

665 # Whereas datasets added in this transaction should not 

666 with self.assertRaises(TransactionTestError): 

667 with datastore.transaction(): 

668 for ref, metrics in fail: 

669 datastore.put(metrics, ref) 

670 raise TransactionTestError("This should propagate out of the context manager") 

671 # Check for datasets that should exist 

672 for ref, metrics in succeed: 

673 # Does it exist? 

674 self.assertTrue(datastore.exists(ref)) 

675 # Get 

676 metricsOut = datastore.get(ref, parameters=None) 

677 self.assertEqual(metrics, metricsOut) 

678 # URI 

679 uri = datastore.getURI(ref) 

680 self.assertEqual(uri.scheme, self.uriScheme) 

681 # Check for datasets that should not exist 

682 for ref, _ in fail: 

683 # These should raise 

684 with self.assertRaises(FileNotFoundError): 

685 # non-existing file 

686 datastore.get(ref) 

687 with self.assertRaises(FileNotFoundError): 

688 datastore.getURI(ref) 

689 

690 def testNestedTransaction(self) -> None: 

691 datastore = self.makeDatastore() 

692 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

693 dimensions = self.universe.conform(("visit", "physical_filter")) 

694 metrics = makeExampleMetrics() 

695 

696 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101} 

697 refBefore = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

698 datastore.put(metrics, refBefore) 

699 with self.assertRaises(TransactionTestError): 

700 with datastore.transaction(): 

701 dataId = { 

702 "instrument": "dummy", 

703 "visit": 1, 

704 "physical_filter": "V", 

705 "band": "v", 

706 "day_obs": 20250101, 

707 } 

708 refOuter = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

709 datastore.put(metrics, refOuter) 

710 with datastore.transaction(): 

711 dataId = { 

712 "instrument": "dummy", 

713 "visit": 2, 

714 "physical_filter": "V", 

715 "band": "v", 

716 "day_obs": 20250101, 

717 } 

718 refInner = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

719 datastore.put(metrics, refInner) 

720 # All datasets should exist 

721 for ref in (refBefore, refOuter, refInner): 

722 metricsOut = datastore.get(ref, parameters=None) 

723 self.assertEqual(metrics, metricsOut) 

724 raise TransactionTestError("This should roll back the transaction") 

725 # Dataset(s) inserted before the transaction should still exist 

726 metricsOut = datastore.get(refBefore, parameters=None) 

727 self.assertEqual(metrics, metricsOut) 

728 # But all datasets inserted during the (rolled back) transaction 

729 # should be gone 

730 with self.assertRaises(FileNotFoundError): 

731 datastore.get(refOuter) 

732 with self.assertRaises(FileNotFoundError): 

733 datastore.get(refInner) 

734 

735 def _prepareIngestTest(self) -> tuple[MetricsExample, DatasetRef]: 

736 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

737 dimensions = self.universe.conform(("visit", "physical_filter")) 

738 metrics = makeExampleMetrics() 

739 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101} 

740 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

741 return metrics, ref 

742 

743 def runIngestTest(self, func: Callable[[MetricsExample, str, DatasetRef], None]) -> None: 

744 metrics, ref = self._prepareIngestTest() 

745 # The file will be deleted after the test. 

746 # For symlink tests this leads to a situation where the datastore 

747 # points to a file that does not exist. This will make os.path.exist 

748 # return False but then the new symlink will fail with 

749 # FileExistsError later in the code so the test still passes. 

750 with _temp_yaml_file(metrics._asdict()) as path: 

751 func(metrics, path, ref) 

752 

753 def testIngestNoTransfer(self) -> None: 

754 """Test ingesting existing files with no transfer.""" 

755 for mode in (None, "auto"): 

756 # Some datastores have auto but can't do in place transfer 

757 if mode == "auto" and "auto" in self.ingestTransferModes and not self.canIngestNoTransferAuto: 

758 continue 

759 

760 with self.subTest(mode=mode): 

761 datastore = self.makeDatastore() 

762 

763 def succeed( 

764 obj: MetricsExample, 

765 path: str, 

766 ref: DatasetRef, 

767 mode: str | None = mode, 

768 datastore: Datastore = datastore, 

769 ) -> None: 

770 """Ingest a file already in the datastore root.""" 

771 # first move it into the root, and adjust the path 

772 # accordingly. 

773 # In the case of a ChainedDatastore, we have multiple 

774 # roots, all of which will accept the file, so we 

775 # have to copy it into all the roots. 

776 relative_path = None 

777 for root in datastore.roots.values(): 

778 if root is not None: 

779 copied_path = shutil.copy(path, root.ospath) 

780 relative_path = os.path.relpath(copied_path, start=root.ospath) 

781 assert relative_path is not None, ( 

782 "Running a FileDatastore test on a Datastore instance without any roots" 

783 ) 

784 datastore.ingest(FileDataset(path=relative_path, refs=ref), transfer=mode) 

785 self.assertEqual(obj, datastore.get(ref)) 

786 

787 def failInputDoesNotExist( 

788 obj: MetricsExample, 

789 path: str, 

790 ref: DatasetRef, 

791 mode: str | None = mode, 

792 datastore: Datastore = datastore, 

793 ) -> None: 

794 """Can't ingest files if we're given a bad path.""" 

795 with self.assertRaises(FileNotFoundError): 

796 datastore.ingest( 

797 FileDataset(path="this-file-does-not-exist.yaml", refs=ref), transfer=mode 

798 ) 

799 self.assertFalse(datastore.exists(ref)) 

800 

801 def failOutsideRoot( 

802 obj: MetricsExample, 

803 path: str, 

804 ref: DatasetRef, 

805 mode: str | None = mode, 

806 datastore: Datastore = datastore, 

807 ) -> None: 

808 """Can't ingest files outside of datastore root unless 

809 auto. 

810 """ 

811 if mode == "auto": 

812 datastore.ingest(FileDataset(path=os.path.abspath(path), refs=ref), transfer=mode) 

813 self.assertTrue(datastore.exists(ref)) 

814 else: 

815 with self.assertRaises(RuntimeError): 

816 datastore.ingest(FileDataset(path=os.path.abspath(path), refs=ref), transfer=mode) 

817 self.assertFalse(datastore.exists(ref)) 

818 

819 def failNotImplemented( 

820 obj: MetricsExample, 

821 path: str, 

822 ref: DatasetRef, 

823 mode: str | None = mode, 

824 datastore: Datastore = datastore, 

825 ) -> None: 

826 with self.assertRaises(NotImplementedError): 

827 datastore.ingest(FileDataset(path=path, refs=ref), transfer=mode) 

828 

829 if mode in self.ingestTransferModes: 

830 self.runIngestTest(failOutsideRoot) 

831 self.runIngestTest(failInputDoesNotExist) 

832 self.runIngestTest(succeed) 

833 else: 

834 self.runIngestTest(failNotImplemented) 

835 

836 def testIngestTransfer(self) -> None: 

837 """Test ingesting existing files after transferring them.""" 

838 for mode in ("copy", "move", "link", "hardlink", "symlink", "relsymlink", "auto"): 

839 with self.subTest(mode=mode): 

840 datastore = self.makeDatastore(mode) 

841 

842 def succeed( 

843 obj: MetricsExample, 

844 path: str, 

845 ref: DatasetRef, 

846 mode: str | None = mode, 

847 datastore: Datastore = datastore, 

848 ) -> None: 

849 """Ingest a file by transferring it to the template 

850 location. 

851 """ 

852 datastore.ingest(FileDataset(path=os.path.abspath(path), refs=ref), transfer=mode) 

853 self.assertEqual(obj, datastore.get(ref)) 

854 file_exists = os.path.exists(path) 

855 if mode == "move": 

856 self.assertFalse(file_exists) 

857 else: 

858 self.assertTrue(file_exists) 

859 

860 def failInputDoesNotExist( 

861 obj: MetricsExample, 

862 path: str, 

863 ref: DatasetRef, 

864 mode: str | None = mode, 

865 datastore: Datastore = datastore, 

866 ) -> None: 

867 """Can't ingest files if we're given a bad path.""" 

868 with self.assertRaises(FileNotFoundError): 

869 # Ensure the file does not look like it is in 

870 # datastore for auto mode 

871 datastore.ingest( 

872 FileDataset(path="../this-file-does-not-exist.yaml", refs=ref), transfer=mode 

873 ) 

874 self.assertFalse(datastore.exists(ref), f"Checking not in datastore using mode {mode}") 

875 

876 def failNotImplemented( 

877 obj: MetricsExample, 

878 path: str, 

879 ref: DatasetRef, 

880 mode: str | None = mode, 

881 datastore: Datastore = datastore, 

882 ) -> None: 

883 with self.assertRaises(NotImplementedError): 

884 datastore.ingest(FileDataset(path=os.path.abspath(path), refs=ref), transfer=mode) 

885 

886 if mode in self.ingestTransferModes: 

887 self.runIngestTest(failInputDoesNotExist) 

888 self.runIngestTest(succeed) 

889 else: 

890 self.runIngestTest(failNotImplemented) 

891 

892 def testIngestSymlinkOfSymlink(self) -> None: 

893 """Special test for symlink to a symlink ingest""" 

894 metrics, ref = self._prepareIngestTest() 

895 # The aim of this test is to create a dataset on disk, then 

896 # create a symlink to it and finally ingest the symlink such that 

897 # the symlink in the datastore points to the original dataset. 

898 for mode in ("symlink", "relsymlink"): 

899 if mode not in self.ingestTransferModes: 

900 continue 

901 

902 print(f"Trying mode {mode}") 

903 with _temp_yaml_file(metrics._asdict()) as realpath: 

904 with tempfile.TemporaryDirectory() as tmpdir: 

905 sympath = os.path.join(tmpdir, "symlink.yaml") 

906 os.symlink(os.path.realpath(realpath), sympath) 

907 

908 datastore = self.makeDatastore() 

909 datastore.ingest(FileDataset(path=os.path.abspath(sympath), refs=ref), transfer=mode) 

910 

911 uri = datastore.getURI(ref) 

912 self.assertTrue(uri.isLocal, f"Check {uri.scheme}") 

913 self.assertTrue(os.path.islink(uri.ospath), f"Check {uri} is a symlink") 

914 

915 linkTarget = os.readlink(uri.ospath) 

916 if mode == "relsymlink": 

917 self.assertFalse(os.path.isabs(linkTarget)) 

918 else: 

919 self.assertTrue(os.path.samefile(linkTarget, realpath)) 

920 

921 # Check that we can get the dataset back regardless of mode 

922 metric2 = datastore.get(ref) 

923 self.assertEqual(metric2, metrics) 

924 

925 # Cleanup the file for next time round loop 

926 # since it will get the same file name in store 

927 datastore.remove(ref) 

928 

929 def _populate_export_datastore(self, name: str) -> tuple[Datastore, list[DatasetRef]]: 

930 datastore = self.makeDatastore(name) 

931 

932 # For now only the FileDatastore can be used for this test. 

933 # ChainedDatastore that only includes InMemoryDatastores have to be 

934 # skipped as well. 

935 for name in datastore.names: 

936 if not name.startswith("InMemoryDatastore"): 

937 break 

938 else: 

939 raise unittest.SkipTest("in-memory datastore does not support record export/import") 

940 

941 metrics = makeExampleMetrics() 

942 dimensions = self.universe.conform(("visit", "physical_filter")) 

943 sc = self.storageClassFactory.getStorageClass("StructuredData") 

944 

945 refs = [] 

946 for visit in (2048, 2049, 2050): 

947 dataId = { 

948 "instrument": "dummy", 

949 "visit": visit, 

950 "physical_filter": "Uprime", 

951 "band": "u", 

952 "day_obs": 20250101, 

953 } 

954 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

955 datastore.put(metrics, ref) 

956 refs.append(ref) 

957 return datastore, refs 

958 

959 def testExportImportRecords(self) -> None: 

960 """Test for export_records and import_records methods.""" 

961 datastore, refs = self._populate_export_datastore("test_datastore") 

962 for exported_refs in (refs, refs[1:]): 

963 n_refs = len(exported_refs) 

964 records = datastore.export_records(exported_refs) 

965 self.assertGreater(len(records), 0) 

966 self.assertTrue(set(records.keys()) <= set(datastore.names)) 

967 # In a ChainedDatastore each FileDatastore will have a complete set 

968 for datastore_name in records: 

969 record_data = records[datastore_name] 

970 self.assertEqual(len(record_data.records), n_refs) 

971 

972 # Check that subsetting works, include non-existing dataset ID. 

973 dataset_ids = {exported_refs[0].id, uuid.uuid4()} 

974 subset = record_data.subset(dataset_ids) 

975 assert subset is not None 

976 self.assertEqual(len(subset.records), 1) 

977 subset = record_data.subset({uuid.uuid4()}) 

978 self.assertIsNone(subset) 

979 

980 # Use the same datastore name to import relative path. 

981 datastore2 = self.makeDatastore("test_datastore") 

982 

983 records = datastore.export_records(refs[1:]) 

984 datastore2.import_records(records) 

985 

986 with self.assertRaises(FileNotFoundError): 

987 data = datastore2.get(refs[0]) 

988 data = datastore2.get(refs[1]) 

989 self.assertIsNotNone(data) 

990 data = datastore2.get(refs[2]) 

991 self.assertIsNotNone(data) 

992 

993 def testExportPredictedRecords(self): 

994 if self.isEphemeral: 

995 raise unittest.SkipTest("in-memory datastore does not support record export/import") 

996 sc = self.storageClassFactory.getStorageClass("ThingOne") 

997 dimensions = self.universe.conform(("visit", "physical_filter")) 

998 dataId = { 

999 "instrument": "dummy", 

1000 "visit": 52, 

1001 "physical_filter": "V", 

1002 "band": "v", 

1003 "day_obs": 20250101, 

1004 } 

1005 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

1006 

1007 datastore = self.makeDatastore("test_datastore") 

1008 names = {n for n in datastore.names if not n.startswith("InMemory")} 

1009 records = datastore.export_predicted_records([ref]) 

1010 

1011 # Expect predicted records from all datastores. 

1012 self.assertEqual(set(records.keys()), names) 

1013 

1014 for record_data in records.values(): 

1015 self.assertEqual(len(record_data.records), 1) 

1016 

1017 def testExport(self) -> None: 

1018 datastore, refs = self._populate_export_datastore("test_datastore") 

1019 

1020 datasets = list(datastore.export(refs)) 

1021 self.assertEqual(len(datasets), 3) 

1022 

1023 for transfer in (None, "auto"): 

1024 # Both will default to None 

1025 datasets = list(datastore.export(refs, transfer=transfer)) 

1026 self.assertEqual(len(datasets), 3) 

1027 

1028 with self.assertRaises(TypeError): 

1029 list(datastore.export(refs, transfer="copy")) 

1030 

1031 with self.assertRaises(TypeError): 

1032 list(datastore.export(refs, directory="exportDir", transfer="move")) 

1033 

1034 # Create a new ref that is not known to the datastore and try to 

1035 # export it. 

1036 sc = self.storageClassFactory.getStorageClass("ThingOne") 

1037 dimensions = self.universe.conform(("visit", "physical_filter")) 

1038 dataId = { 

1039 "instrument": "dummy", 

1040 "visit": 52, 

1041 "physical_filter": "V", 

1042 "band": "v", 

1043 "day_obs": 20250101, 

1044 } 

1045 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

1046 with self.assertRaises(FileNotFoundError): 

1047 list(datastore.export(refs + [ref], transfer=None)) 

1048 

1049 def test_pydantic_dict_storage_class_conversions(self) -> None: 

1050 """Test converting a dataset stored as a pydantic model into a dict on 

1051 read. 

1052 """ 

1053 datastore = self.makeDatastore() 

1054 store_as_model = self.makeDatasetRef( 

1055 "store_as_model", 

1056 dimensions=self.universe.empty, 

1057 storageClass="DictConvertibleModel", 

1058 dataId=DataCoordinate.make_empty(self.universe), 

1059 ) 

1060 content = {"a": "one", "b": "two"} 

1061 model = DictConvertibleModel.from_dict(content, extra="original content") 

1062 datastore.put(model, store_as_model) 

1063 retrieved_model = datastore.get(store_as_model) 

1064 self.assertEqual(retrieved_model, model) 

1065 loaded = datastore.get(store_as_model.overrideStorageClass("NativeDictForConvertibleModel")) 

1066 self.assertEqual(type(loaded), dict) 

1067 self.assertEqual(loaded, content) 

1068 

1069 def test_simple_class_put_get(self) -> None: 

1070 """Test that we can put and get a simple class with dict() 

1071 constructor. 

1072 """ 

1073 datastore = self.makeDatastore() 

1074 data = MetricsExample(summary={"a": 1}, data=[1, 2, 3], output={"b": 2}) 

1075 self._assert_different_puts(datastore, "MetricsExample", data) 

1076 

1077 def test_dataclass_put_get(self) -> None: 

1078 """Test that we can put and get a simple dataclass.""" 

1079 datastore = self.makeDatastore() 

1080 data = MetricsExampleDataclass(summary={"a": 1}, data=[1, 2, 3], output={"b": 2}) 

1081 self._assert_different_puts(datastore, "MetricsExampleDataclass", data) 

1082 

1083 def test_pydantic_put_get(self) -> None: 

1084 """Test that we can put and get a simple Pydantic model.""" 

1085 datastore = self.makeDatastore() 

1086 data = MetricsExampleModel(summary={"a": 1}, data=[1, 2, 3], output={"b": 2}) 

1087 self._assert_different_puts(datastore, "MetricsExampleModel", data) 

1088 

1089 def test_tuple_put_get(self) -> None: 

1090 """Test that we can put and get a tuple.""" 

1091 datastore = self.makeDatastore() 

1092 data = ("a", "b", 1) 

1093 self._assert_different_puts(datastore, "TupleExample", data) 

1094 

1095 def _assert_different_puts(self, datastore: Datastore, storageClass_root: str, data: Any) -> None: 

1096 refs = { 

1097 x: self.makeDatasetRef( 

1098 f"stora_as_{x}", 

1099 dimensions=self.universe.empty, 

1100 storageClass=f"{storageClass_root}{x}", 

1101 dataId=DataCoordinate.make_empty(self.universe), 

1102 ) 

1103 for x in ["A", "B"] 

1104 } 

1105 

1106 for ref in refs.values(): 

1107 datastore.put(data, ref) 

1108 

1109 self.assertEqual(datastore.get(refs["A"]), datastore.get(refs["B"])) 

1110 

1111 

1112class PosixDatastoreTestCase(DatastoreTests, unittest.TestCase): 

1113 """PosixDatastore specialization""" 

1114 

1115 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

1116 uriScheme = "file" 

1117 canIngestNoTransferAuto = True 

1118 ingestTransferModes = (None, "copy", "move", "link", "hardlink", "symlink", "relsymlink", "auto") 

1119 isEphemeral = False 

1120 rootKeys = ("root",) 

1121 validationCanFail = True 

1122 

1123 def setUp(self) -> None: 

1124 # The call to os.path.realpath is necessary because Mac temporary files 

1125 # can end up in either /private/var/folders or /var/folders, which 

1126 # refer to the same location but don't appear to. 

1127 # This matters for "relsymlink" transfer mode, because it needs to be 

1128 # able to read the file through a relative symlink, but some of the 

1129 # intermediate directories are not traversable if you try to get from a 

1130 # tempfile in /var/folders to one in /private/var/folders via a 

1131 # relative path. 

1132 self.root = os.path.realpath(self.enterContext(tempfile.TemporaryDirectory())) 

1133 super().setUp() 

1134 

1135 def testAtomicWrite(self) -> None: 

1136 """Test that we write to a temporary and then rename""" 

1137 datastore = self.makeDatastore() 

1138 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

1139 dimensions = self.universe.conform(("visit", "physical_filter")) 

1140 metrics = makeExampleMetrics() 

1141 

1142 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101} 

1143 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

1144 

1145 with self.assertLogs("lsst.resources", "DEBUG") as cm: 

1146 datastore.put(metrics, ref) 

1147 move_logs = [ll for ll in cm.output if "transfer=" in ll] 

1148 self.assertIn("transfer=move", move_logs[0]) 

1149 

1150 # And the transfer should be file to file. 

1151 self.assertEqual(move_logs[0].count("file://"), 2) 

1152 

1153 def testCanNotDeterminePutFormatterLocation(self) -> None: 

1154 """Verify that the expected exception is raised if the FileDatastore 

1155 can not determine the put formatter location. 

1156 """ 

1157 _ = makeExampleMetrics() 

1158 datastore = self.makeDatastore() 

1159 

1160 # Create multiple storage classes for testing different formulations 

1161 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

1162 

1163 sccomp = StorageClass("Dummy") 

1164 compositeStorageClass = StorageClass( 

1165 "StructuredComposite", components={"dummy": sccomp, "dummy2": sccomp} 

1166 ) 

1167 

1168 dimensions = self.universe.conform(("visit", "physical_filter")) 

1169 dataId = { 

1170 "instrument": "dummy", 

1171 "visit": 52, 

1172 "physical_filter": "V", 

1173 "band": "v", 

1174 "day_obs": 20250101, 

1175 } 

1176 

1177 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

1178 compRef = self.makeDatasetRef("metric", dimensions, compositeStorageClass, dataId) 

1179 

1180 def raiser(ref: DatasetRef) -> None: 

1181 raise DatasetTypeNotSupportedError() 

1182 

1183 with unittest.mock.patch.object( 

1184 lsst.daf.butler.datastores.fileDatastore.FileDatastore, 

1185 "_determine_put_formatter_location", 

1186 side_effect=raiser, 

1187 ): 

1188 # verify the non-composite ref execution path: 

1189 with self.assertRaises(DatasetTypeNotSupportedError): 

1190 datastore.getURIs(ref, predict=True) 

1191 

1192 # verify the composite-ref execution path: 

1193 with self.assertRaises(DatasetTypeNotSupportedError): 

1194 datastore.getURIs(compRef, predict=True) 

1195 

1196 def test_roots(self): 

1197 datastore = self.makeDatastore() 

1198 

1199 self.assertEqual(set(datastore.names), set(datastore.roots.keys())) 

1200 for root in datastore.roots.values(): 

1201 if root is not None: 

1202 self.assertTrue(root.exists()) 

1203 

1204 def test_prepare_get_for_external_client(self): 

1205 datastore = self.makeDatastore() 

1206 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

1207 dimensions = self.universe.conform(("visit", "physical_filter")) 

1208 dataId = {"instrument": "dummy", "visit": 52, "physical_filter": "V", "band": "v"} 

1209 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

1210 # Most of the coverage for this function is in test_server.py, 

1211 # because it requires a file backend that supports URL signing. 

1212 self.assertIsNone(datastore.prepare_get_for_external_client(ref)) 

1213 

1214 

1215class PosixDatastoreNoChecksumsTestCase(PosixDatastoreTestCase): 

1216 """Posix datastore tests but with checksums disabled.""" 

1217 

1218 configFile = os.path.join(TESTDIR, "config/basic/posixDatastoreNoChecksums.yaml") 

1219 

1220 def testChecksum(self) -> None: 

1221 """Ensure that checksums have not been calculated.""" 

1222 datastore = self.makeDatastore() 

1223 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

1224 dimensions = self.universe.conform(("visit", "physical_filter")) 

1225 metrics = makeExampleMetrics() 

1226 

1227 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101} 

1228 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

1229 

1230 # Configuration should have disabled checksum calculation 

1231 datastore.put(metrics, ref) 

1232 infos = datastore.getStoredItemsInfo(ref) 

1233 self.assertIsNone(infos[0].checksum) 

1234 

1235 # Remove put back but with checksums enabled explicitly 

1236 datastore.remove(ref) 

1237 datastore.useChecksum = True 

1238 datastore.put(metrics, ref) 

1239 

1240 infos = datastore.getStoredItemsInfo(ref) 

1241 self.assertIsNotNone(infos[0].checksum) 

1242 

1243 def test_repeat_ingest(self): 

1244 """Test that repeatedly ingesting the same file in direct mode 

1245 is allowed. 

1246 

1247 Test can only run with FileDatastore since that is the only one 

1248 supporting "direct" ingest. 

1249 """ 

1250 metrics, v4ref = self._prepareIngestTest() 

1251 datastore = self.makeDatastore() 

1252 v5ref = DatasetRef( 

1253 v4ref.datasetType, v4ref.dataId, v4ref.run, id_generation_mode=DatasetIdGenEnum.DATAID_TYPE_RUN 

1254 ) 

1255 

1256 with _temp_yaml_file(metrics._asdict()) as path: 

1257 datastore.ingest(FileDataset(path=path, refs=v4ref), transfer="direct") 

1258 

1259 # This will fail because the ref is using UUIDv4. 

1260 with self.assertRaises(RuntimeError): 

1261 datastore.ingest(FileDataset(path=path, refs=v4ref), transfer="direct") 

1262 

1263 # UUIDv5 can be repeatedly ingested in direct mode. 

1264 datastore.ingest(FileDataset(path=path, refs=v5ref), transfer="direct") 

1265 datastore.ingest(FileDataset(path=path, refs=v5ref), transfer="direct") 

1266 

1267 with self.assertRaises(RuntimeError): 

1268 datastore.ingest(FileDataset(path=path, refs=v5ref), transfer="copy") 

1269 

1270 

1271class TrashDatastoreTestCase(PosixDatastoreTestCase): 

1272 """Restrict trash test to FileDatastore.""" 

1273 

1274 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

1275 

1276 def testTrash(self) -> None: 

1277 datastore, *refs = self.prepDeleteTest(n_refs=10) 

1278 

1279 # Trash one of them. 

1280 ref = refs.pop() 

1281 uri = datastore.getURI(ref) 

1282 datastore.trash(ref) 

1283 self.assertTrue(uri.exists(), uri) # Not deleted yet 

1284 datastore.emptyTrash() 

1285 self.assertFalse(uri.exists(), uri) 

1286 

1287 # Trash it again should be fine. 

1288 datastore.trash(ref) 

1289 

1290 # Trash multiple items at once. 

1291 subset = [refs.pop(), refs.pop()] 

1292 datastore.trash(subset) 

1293 datastore.emptyTrash() 

1294 

1295 # Remove a record and trash should do nothing. 

1296 # This is execution butler scenario. 

1297 ref = refs.pop() 

1298 uri = datastore.getURI(ref) 

1299 datastore._table.delete(["dataset_id"], {"dataset_id": ref.id}) 

1300 self.assertTrue(uri.exists()) 

1301 datastore.trash(ref) 

1302 datastore.emptyTrash() 

1303 self.assertTrue(uri.exists()) 

1304 

1305 # Switch on trust and it should delete the file. 

1306 datastore.trustGetRequest = True 

1307 datastore.trash([ref]) 

1308 self.assertFalse(uri.exists()) 

1309 

1310 # Remove multiples at once in trust mode. 

1311 subset = [refs.pop() for i in range(3)] 

1312 datastore.trash(subset) 

1313 datastore.trash(refs.pop()) # Check that a single ref can trash 

1314 

1315 def test_empty_trash(self) -> None: 

1316 """Test parameters and return value for empty trash.""" 

1317 datastore, *refs = self.prepDeleteTest(n_refs=10) 

1318 

1319 # Trash one of them. 

1320 ref = refs.pop() 

1321 uri = datastore.getURI(ref) 

1322 datastore.trash(ref) 

1323 self.assertTrue(uri.exists(), uri) # Not deleted yet 

1324 

1325 # Empty trash but with a list of refs that does not include the 

1326 # one in the trash table. 

1327 removed = datastore.emptyTrash(refs=refs) 

1328 self.assertEqual(len(removed), 0) 

1329 self.assertTrue(uri.exists(), uri) 

1330 

1331 # Empty the entire trash but in dry_run mode. 

1332 removed = datastore.emptyTrash(dry_run=True) 

1333 self.assertEqual(len(removed), 1) 

1334 self.assertEqual(removed.pop(), uri) 

1335 self.assertTrue(uri.exists(), uri) 

1336 

1337 # Empty the trash specifying the actual ref. 

1338 removed = datastore.emptyTrash(refs=[ref]) 

1339 self.assertEqual(len(removed), 1) 

1340 self.assertEqual(removed.pop(), uri) 

1341 self.assertFalse(uri.exists(), uri) 

1342 

1343 # Trash everything and empty. 

1344 datastore.trash(refs) 

1345 removed = datastore.emptyTrash(dry_run=True) 

1346 for u in removed: 

1347 self.assertTrue(u.exists()) 

1348 removed = datastore.emptyTrash() 

1349 for u in removed: 

1350 self.assertFalse(u.exists()) 

1351 

1352 

1353class CleanupPosixDatastoreTestCase(DatastoreTestsBase, unittest.TestCase): 

1354 """Test datastore cleans up on failure.""" 

1355 

1356 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

1357 

1358 def setUp(self) -> None: 

1359 # Override the working directory before calling the base class 

1360 self.root = tempfile.mkdtemp() 

1361 super().setUp() 

1362 

1363 def testCleanup(self) -> None: 

1364 """Test that a failed formatter write does cleanup a partial file.""" 

1365 metrics = makeExampleMetrics() 

1366 datastore = self.makeDatastore() 

1367 

1368 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

1369 

1370 dimensions = self.universe.conform(("visit", "physical_filter")) 

1371 dataId = { 

1372 "instrument": "dummy", 

1373 "visit": 52, 

1374 "physical_filter": "V", 

1375 "band": "v", 

1376 "day_obs": 20250101, 

1377 } 

1378 

1379 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

1380 

1381 # Determine where the file will end up (we assume Formatters use 

1382 # the same file extension) 

1383 expectedUri = datastore.getURI(ref, predict=True) 

1384 self.assertEqual(expectedUri.fragment, "predicted") 

1385 

1386 self.assertEqual(expectedUri.getExtension(), ".yaml", f"Is there a file extension in {expectedUri}") 

1387 

1388 # Try formatter that fails and formatter that fails and leaves 

1389 # a file behind 

1390 for formatter in (BadWriteFormatter, BadNoWriteFormatter): 

1391 with self.subTest(formatter=str(formatter)): 

1392 # Monkey patch the formatter 

1393 datastore.formatterFactory.registerFormatter(ref.datasetType, formatter, overwrite=True) 

1394 

1395 # Try to put the dataset, it should fail 

1396 with self.assertRaises(RuntimeError): 

1397 datastore.put(metrics, ref) 

1398 

1399 # Check that there is no file on disk 

1400 self.assertFalse(expectedUri.exists(), f"Check for existence of {expectedUri}") 

1401 

1402 # Check that there is a directory 

1403 dir = expectedUri.dirname() 

1404 self.assertTrue(dir.exists(), f"Check for existence of directory {dir}") 

1405 

1406 # Force YamlFormatter and check that this time a file is written 

1407 datastore.formatterFactory.registerFormatter(ref.datasetType, YamlFormatter, overwrite=True) 

1408 datastore.put(metrics, ref) 

1409 self.assertTrue(expectedUri.exists(), f"Check for existence of {expectedUri}") 

1410 datastore.remove(ref) 

1411 self.assertFalse(expectedUri.exists(), f"Check for existence of now removed {expectedUri}") 

1412 

1413 

1414class InMemoryDatastoreTestCase(DatastoreTests, unittest.TestCase): 

1415 """PosixDatastore specialization""" 

1416 

1417 configFile = os.path.join(TESTDIR, "config/basic/inMemoryDatastore.yaml") 

1418 uriScheme = "mem" 

1419 hasUnsupportedPut = False 

1420 ingestTransferModes = () 

1421 isEphemeral = True 

1422 rootKeys = None 

1423 validationCanFail = False 

1424 

1425 

1426class ChainedDatastoreTestCase(PosixDatastoreTestCase): 

1427 """ChainedDatastore specialization using a POSIXDatastore""" 

1428 

1429 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastore.yaml") 

1430 hasUnsupportedPut = False 

1431 canIngestNoTransferAuto = False 

1432 ingestTransferModes = (None, "copy", "move", "hardlink", "symlink", "relsymlink", "link", "auto") 

1433 isEphemeral = False 

1434 rootKeys = (".datastores.1.root", ".datastores.2.root") 

1435 validationCanFail = True 

1436 

1437 

1438class ChainedDatastoreMemoryTestCase(InMemoryDatastoreTestCase): 

1439 """ChainedDatastore specialization using all InMemoryDatastore""" 

1440 

1441 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastore2.yaml") 

1442 validationCanFail = False 

1443 isEphemeral = True 

1444 

1445 

1446class DatastoreConstraintsTests(DatastoreTestsBase): 

1447 """Basic tests of constraints model of Datastores.""" 

1448 

1449 def testConstraints(self) -> None: 

1450 """Test constraints model. Assumes that each test class has the 

1451 same constraints. 

1452 """ 

1453 metrics = makeExampleMetrics() 

1454 datastore = self.makeDatastore() 

1455 

1456 sc1 = self.storageClassFactory.getStorageClass("StructuredData") 

1457 sc2 = self.storageClassFactory.getStorageClass("StructuredDataJson") 

1458 dimensions = self.universe.conform(("visit", "physical_filter", "instrument")) 

1459 dataId = { 

1460 "visit": 52, 

1461 "physical_filter": "V", 

1462 "band": "v", 

1463 "instrument": "DummyCamComp", 

1464 "day_obs": 20250101, 

1465 } 

1466 

1467 # Write empty file suitable for ingest check (JSON and YAML variants) 

1468 testfile_y = tempfile.NamedTemporaryFile(suffix=".yaml") 

1469 testfile_j = tempfile.NamedTemporaryFile(suffix=".json") 

1470 for datasetTypeName, sc, accepted in ( 

1471 ("metric", sc1, True), 

1472 ("metric5", sc1, False), 

1473 ("metric33", sc1, True), 

1474 ("metric5", sc2, True), 

1475 ): 

1476 # Choose different temp file depending on StorageClass 

1477 testfile = testfile_j if sc.name.endswith("Json") else testfile_y 

1478 

1479 with self.subTest(datasetTypeName=datasetTypeName, storageClass=sc.name, file=testfile.name): 

1480 ref = self.makeDatasetRef(datasetTypeName, dimensions, sc, dataId) 

1481 if accepted: 

1482 datastore.put(metrics, ref) 

1483 self.assertTrue(datastore.exists(ref)) 

1484 datastore.remove(ref) 

1485 

1486 # Try ingest 

1487 if self.canIngest: 

1488 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link") 

1489 self.assertTrue(datastore.exists(ref)) 

1490 datastore.remove(ref) 

1491 else: 

1492 with self.assertRaises(DatasetTypeNotSupportedError): 

1493 datastore.put(metrics, ref) 

1494 self.assertFalse(datastore.exists(ref)) 

1495 

1496 # Again with ingest 

1497 if self.canIngest: 

1498 with self.assertRaises(DatasetTypeNotSupportedError): 

1499 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link") 

1500 self.assertFalse(datastore.exists(ref)) 

1501 

1502 

1503class PosixDatastoreConstraintsTestCase(DatastoreConstraintsTests, unittest.TestCase): 

1504 """PosixDatastore specialization""" 

1505 

1506 configFile = os.path.join(TESTDIR, "config/basic/posixDatastoreP.yaml") 

1507 canIngest = True 

1508 

1509 def setUp(self) -> None: 

1510 # Override the working directory before calling the base class 

1511 self.root = tempfile.mkdtemp() 

1512 super().setUp() 

1513 

1514 

1515class InMemoryDatastoreConstraintsTestCase(DatastoreConstraintsTests, unittest.TestCase): 

1516 """InMemoryDatastore specialization.""" 

1517 

1518 configFile = os.path.join(TESTDIR, "config/basic/inMemoryDatastoreP.yaml") 

1519 canIngest = False 

1520 

1521 

1522class ChainedDatastoreConstraintsNativeTestCase(PosixDatastoreConstraintsTestCase): 

1523 """ChainedDatastore specialization using a POSIXDatastore and constraints 

1524 at the ChainedDatstore. 

1525 """ 

1526 

1527 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastorePa.yaml") 

1528 

1529 

1530class ChainedDatastoreConstraintsTestCase(PosixDatastoreConstraintsTestCase): 

1531 """ChainedDatastore specialization using a POSIXDatastore.""" 

1532 

1533 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastoreP.yaml") 

1534 

1535 

1536class ChainedDatastoreMemoryConstraintsTestCase(InMemoryDatastoreConstraintsTestCase): 

1537 """ChainedDatastore specialization using all InMemoryDatastore.""" 

1538 

1539 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastore2P.yaml") 

1540 canIngest = False 

1541 

1542 

1543class ChainedDatastorePerStoreConstraintsTests(DatastoreTestsBase, unittest.TestCase): 

1544 """Test that a chained datastore can control constraints per-datastore 

1545 even if child datastore would accept. 

1546 """ 

1547 

1548 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastorePb.yaml") 

1549 

1550 def setUp(self) -> None: 

1551 # Override the working directory before calling the base class 

1552 self.root = tempfile.mkdtemp() 

1553 super().setUp() 

1554 

1555 def testConstraints(self) -> None: 

1556 """Test chained datastore constraints model.""" 

1557 metrics = makeExampleMetrics() 

1558 datastore = self.makeDatastore() 

1559 

1560 sc1 = self.storageClassFactory.getStorageClass("StructuredData") 

1561 sc2 = self.storageClassFactory.getStorageClass("StructuredDataJson") 

1562 dimensions = self.universe.conform(("visit", "physical_filter", "instrument")) 

1563 dataId1 = { 

1564 "visit": 52, 

1565 "physical_filter": "V", 

1566 "band": "v", 

1567 "instrument": "DummyCamComp", 

1568 "day_obs": 20250101, 

1569 } 

1570 dataId2 = {"visit": 52, "physical_filter": "V", "band": "v", "instrument": "HSC", "day_obs": 20250101} 

1571 

1572 # Write empty file suitable for ingest check (JSON and YAML variants) 

1573 testfile_y = tempfile.NamedTemporaryFile(suffix=".yaml") 

1574 testfile_j = tempfile.NamedTemporaryFile(suffix=".json") 

1575 

1576 for typeName, dataId, sc, accept, ingest in ( 

1577 ("metric", dataId1, sc1, (False, True, False), True), 

1578 ("metric5", dataId1, sc1, (False, False, False), False), 

1579 ("metric5", dataId2, sc1, (True, False, False), False), 

1580 ("metric33", dataId2, sc2, (True, True, False), True), 

1581 ("metric5", dataId1, sc2, (False, True, False), True), 

1582 ): 

1583 # Choose different temp file depending on StorageClass 

1584 testfile = testfile_j if sc.name.endswith("Json") else testfile_y 

1585 

1586 with self.subTest(datasetTypeName=typeName, dataId=dataId, sc=sc.name): 

1587 ref = self.makeDatasetRef(typeName, dimensions, sc, dataId) 

1588 if any(accept): 

1589 datastore.put(metrics, ref) 

1590 self.assertTrue(datastore.exists(ref)) 

1591 

1592 # Check each datastore inside the chained datastore 

1593 for childDatastore, expected in zip(datastore.datastores, accept, strict=True): 

1594 self.assertEqual( 

1595 childDatastore.exists(ref), 

1596 expected, 

1597 f"Testing presence of {ref} in datastore {childDatastore.name}", 

1598 ) 

1599 

1600 datastore.remove(ref) 

1601 

1602 # Check that ingest works 

1603 if ingest: 

1604 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link") 

1605 self.assertTrue(datastore.exists(ref)) 

1606 

1607 # Check each datastore inside the chained datastore 

1608 for childDatastore, expected in zip(datastore.datastores, accept, strict=True): 

1609 # Ephemeral datastores means InMemory at the moment 

1610 # and that does not accept ingest of files. 

1611 if childDatastore.isEphemeral: 

1612 expected = False 

1613 self.assertEqual( 

1614 childDatastore.exists(ref), 

1615 expected, 

1616 f"Testing presence of ingested {ref} in datastore {childDatastore.name}", 

1617 ) 

1618 

1619 datastore.remove(ref) 

1620 else: 

1621 with self.assertRaises(DatasetTypeNotSupportedError): 

1622 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link") 

1623 

1624 else: 

1625 with self.assertRaises(DatasetTypeNotSupportedError): 

1626 datastore.put(metrics, ref) 

1627 self.assertFalse(datastore.exists(ref)) 

1628 

1629 # Again with ingest 

1630 with self.assertRaises(DatasetTypeNotSupportedError): 

1631 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link") 

1632 self.assertFalse(datastore.exists(ref)) 

1633 

1634 

1635@unittest.mock.patch.dict(os.environ, {}, clear=True) 

1636class DatastoreCacheTestCase(DatasetTestHelper, unittest.TestCase): 

1637 """Tests for datastore caching infrastructure.""" 

1638 

1639 @classmethod 

1640 def setUpClass(cls) -> None: 

1641 cls.storageClassFactory = StorageClassFactory() 

1642 cls.universe = DimensionUniverse() 

1643 

1644 # Ensure that we load the test storage class definitions. 

1645 scConfigFile = os.path.join(TESTDIR, "config/basic/storageClasses.yaml") 

1646 cls.storageClassFactory.addFromConfig(scConfigFile) 

1647 

1648 def setUp(self) -> None: 

1649 self.id = 0 

1650 

1651 # Create a root that we can use for caching tests. 

1652 self.root = tempfile.mkdtemp() 

1653 

1654 # Create some test dataset refs and associated test files 

1655 sc = self.storageClassFactory.getStorageClass("StructuredDataDict") 

1656 dimensions = self.universe.conform(("visit", "physical_filter")) 

1657 dataId = { 

1658 "instrument": "dummy", 

1659 "visit": 52, 

1660 "physical_filter": "V", 

1661 "band": "v", 

1662 "day_obs": 20250101, 

1663 } 

1664 

1665 # Create list of refs and list of temporary files 

1666 n_datasets = 10 

1667 self.refs = [self.makeDatasetRef(f"metric{n}", dimensions, sc, dataId) for n in range(n_datasets)] 

1668 

1669 root_uri = ResourcePath(self.root, forceDirectory=True) 

1670 self.files = [root_uri.join(f"file{n}.txt") for n in range(n_datasets)] 

1671 

1672 # Create test files. 

1673 for uri in self.files: 

1674 uri.write(b"0123456789") 

1675 

1676 # Create some composite refs with component files. 

1677 sc = self.storageClassFactory.getStorageClass("StructuredData") 

1678 self.composite_refs = [self.makeDatasetRef(f"composite{n}", dimensions, sc, dataId) for n in range(3)] 

1679 self.comp_files = [] 

1680 self.comp_refs = [] 

1681 for n, ref in enumerate(self.composite_refs): 

1682 component_refs = [] 

1683 component_files = [] 

1684 for component in sc.components: 

1685 component_ref = ref.makeComponentRef(component) 

1686 file = root_uri.join(f"composite_file-{n}-{component}.txt") 

1687 component_refs.append(component_ref) 

1688 component_files.append(file) 

1689 file.write(b"9876543210") 

1690 

1691 self.comp_files.append(component_files) 

1692 self.comp_refs.append(component_refs) 

1693 

1694 def tearDown(self) -> None: 

1695 if self.root is not None and os.path.exists(self.root): 

1696 shutil.rmtree(self.root, ignore_errors=True) 

1697 

1698 def _make_cache_manager(self, config_str: str) -> DatastoreCacheManager: 

1699 config = Config.fromYaml(config_str) 

1700 return DatastoreCacheManager(DatastoreCacheManagerConfig(config), universe=self.universe) 

1701 

1702 def testNoCacheDir(self) -> None: 

1703 config_str = """ 

1704cached: 

1705 root: null 

1706 cacheable: 

1707 metric0: true 

1708 """ 

1709 cache_manager = self._make_cache_manager(config_str) 

1710 

1711 # Look inside to check we don't have a cache directory 

1712 self.assertIsNone(cache_manager._cache_directory) 

1713 

1714 self.assertCache(cache_manager) 

1715 

1716 # Test that the cache directory is marked temporary 

1717 self.assertTrue(cache_manager.cache_directory.isTemporary) 

1718 

1719 def testNoCacheDirReversed(self) -> None: 

1720 """Use default caching status and metric1 to false""" 

1721 config_str = """ 

1722cached: 

1723 root: null 

1724 default: true 

1725 cacheable: 

1726 metric1: false 

1727 """ 

1728 cache_manager = self._make_cache_manager(config_str) 

1729 

1730 self.assertCache(cache_manager) 

1731 

1732 def testEnvvarCacheDir(self) -> None: 

1733 config_str = f""" 

1734cached: 

1735 root: '{self.root}' 

1736 cacheable: 

1737 metric0: true 

1738 """ 

1739 

1740 root = ResourcePath(self.root, forceDirectory=True) 

1741 env_dir = root.join("somewhere", forceDirectory=True) 

1742 elsewhere = root.join("elsewhere", forceDirectory=True) 

1743 

1744 # Environment variable should override the config value. 

1745 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_CACHE_DIRECTORY": env_dir.ospath}): 

1746 cache_manager = self._make_cache_manager(config_str) 

1747 self.assertEqual(cache_manager.cache_directory, env_dir) 

1748 

1749 # This environment variable should not override the config value. 

1750 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": env_dir.ospath}): 

1751 cache_manager = self._make_cache_manager(config_str) 

1752 self.assertEqual(cache_manager.cache_directory, root) 

1753 

1754 # No default setting. 

1755 config_str = """ 

1756cached: 

1757 root: null 

1758 default: true 

1759 cacheable: 

1760 metric1: false 

1761 """ 

1762 cache_manager = self._make_cache_manager(config_str) 

1763 

1764 # This environment variable should override the config value. 

1765 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": env_dir.ospath}): 

1766 cache_manager = self._make_cache_manager(config_str) 

1767 self.assertEqual(cache_manager.cache_directory, env_dir) 

1768 

1769 # If both environment variables are set the main (not IF_UNSET) 

1770 # variable should win. 

1771 with unittest.mock.patch.dict( 

1772 os.environ, 

1773 { 

1774 "DAF_BUTLER_CACHE_DIRECTORY": env_dir.ospath, 

1775 "DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": elsewhere.ospath, 

1776 }, 

1777 ): 

1778 cache_manager = self._make_cache_manager(config_str) 

1779 self.assertEqual(cache_manager.cache_directory, env_dir) 

1780 

1781 # Use the API to set the environment variable, making sure that the 

1782 # variable is reset on exit. 

1783 with unittest.mock.patch.dict( 

1784 os.environ, 

1785 {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": ""}, 

1786 ): 

1787 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset() 

1788 self.assertTrue(defined) 

1789 cache_manager = self._make_cache_manager(config_str) 

1790 self.assertEqual(cache_manager.cache_directory, ResourcePath(cache_dir, forceDirectory=True)) 

1791 

1792 # Now create the cache manager ahead of time and set the fallback 

1793 # later. 

1794 cache_manager = self._make_cache_manager(config_str) 

1795 self.assertIsNone(cache_manager._cache_directory) 

1796 with unittest.mock.patch.dict( 

1797 os.environ, 

1798 {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": ""}, 

1799 ): 

1800 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset() 

1801 self.assertTrue(defined) 

1802 self.assertEqual(cache_manager.cache_directory, ResourcePath(cache_dir, forceDirectory=True)) 

1803 

1804 def testExplicitCacheDir(self) -> None: 

1805 config_str = f""" 

1806cached: 

1807 root: '{self.root}' 

1808 cacheable: 

1809 metric0: true 

1810 """ 

1811 cache_manager = self._make_cache_manager(config_str) 

1812 

1813 # Look inside to check we do have a cache directory. 

1814 self.assertEqual(cache_manager.cache_directory, ResourcePath(self.root, forceDirectory=True)) 

1815 

1816 self.assertCache(cache_manager) 

1817 

1818 # Test that the cache directory is not marked temporary 

1819 self.assertFalse(cache_manager.cache_directory.isTemporary) 

1820 

1821 def testUnexpectedFilesInCacheDir(self) -> None: 

1822 """Test for regression of a bug where extraneous files in a cache 

1823 directory would cause all cache lookups to raise an exception. 

1824 """ 

1825 config_str = f""" 

1826cached: 

1827 root: '{self.root}' 

1828 cacheable: 

1829 metric0: true 

1830 """ 

1831 

1832 for filename in ["unexpected.txt", "unexpected", "un_expected", "un_expected.txt"]: 

1833 unexpected_file = os.path.join(self.root, filename) 

1834 with open(unexpected_file, "w") as fh: 

1835 fh.write("test") 

1836 

1837 cache_manager = self._make_cache_manager(config_str) 

1838 cache_manager.scan_cache() 

1839 self.assertCache(cache_manager) 

1840 

1841 def assertCache(self, cache_manager: DatastoreCacheManager) -> None: 

1842 self.assertTrue(cache_manager.should_be_cached(self.refs[0])) 

1843 self.assertFalse(cache_manager.should_be_cached(self.refs[1])) 

1844 

1845 uri = cache_manager.move_to_cache(self.files[0], self.refs[0]) 

1846 self.assertIsInstance(uri, ResourcePath) 

1847 self.assertIsNone(cache_manager.move_to_cache(self.files[1], self.refs[1])) 

1848 

1849 # Check presence in cache using ref and then using file extension. 

1850 self.assertFalse(cache_manager.known_to_cache(self.refs[1])) 

1851 self.assertTrue(cache_manager.known_to_cache(self.refs[0])) 

1852 self.assertFalse(cache_manager.known_to_cache(self.refs[1], self.files[1].getExtension())) 

1853 self.assertTrue(cache_manager.known_to_cache(self.refs[0], self.files[0].getExtension())) 

1854 

1855 # Cached file should no longer exist but uncached file should be 

1856 # unaffected. 

1857 self.assertFalse(self.files[0].exists()) 

1858 self.assertTrue(self.files[1].exists()) 

1859 

1860 # Should find this file and it should be within the cache directory. 

1861 with cache_manager.find_in_cache(self.refs[0], ".txt") as found: 

1862 self.assertTrue(found.exists()) 

1863 self.assertIsNotNone(found.relative_to(cache_manager.cache_directory)) 

1864 

1865 # Should not be able to find these in cache 

1866 with cache_manager.find_in_cache(self.refs[0], ".fits") as found: 

1867 self.assertIsNone(found) 

1868 with cache_manager.find_in_cache(self.refs[1], ".fits") as found: 

1869 self.assertIsNone(found) 

1870 

1871 def testNoCache(self) -> None: 

1872 cache_manager = DatastoreDisabledCacheManager("", universe=self.universe) 

1873 for uri, ref in zip(self.files, self.refs, strict=True): 

1874 self.assertFalse(cache_manager.should_be_cached(ref)) 

1875 self.assertIsNone(cache_manager.move_to_cache(uri, ref)) 

1876 self.assertFalse(cache_manager.known_to_cache(ref)) 

1877 with cache_manager.find_in_cache(ref, ".txt") as found: 

1878 self.assertIsNone(found, msg=f"{cache_manager}") 

1879 

1880 def _expiration_config(self, mode: str, threshold: int) -> str: 

1881 return f""" 

1882cached: 

1883 default: true 

1884 expiry: 

1885 mode: {mode} 

1886 threshold: {threshold} 

1887 cacheable: 

1888 unused: true 

1889 """ 

1890 

1891 def testCacheExpiryFiles(self) -> None: 

1892 threshold = 2 # Keep at least 2 files. 

1893 mode = "files" 

1894 config_str = self._expiration_config(mode, threshold) 

1895 

1896 cache_manager = self._make_cache_manager(config_str) 

1897 

1898 # Check that an empty cache returns unknown for arbitrary ref 

1899 self.assertFalse(cache_manager.known_to_cache(self.refs[0])) 

1900 

1901 # Should end with datasets: 2, 3, 4 

1902 self.assertExpiration(cache_manager, 5, threshold + 1) 

1903 self.assertIn(f"{mode}={threshold}", str(cache_manager)) 

1904 

1905 # Check that we will not expire a file that is actively in use. 

1906 with cache_manager.find_in_cache(self.refs[2], ".txt") as found: 

1907 self.assertIsNotNone(found) 

1908 

1909 # Trigger cache expiration that should remove the file 

1910 # we just retrieved. Should now have: 3, 4, 5 

1911 cached = cache_manager.move_to_cache(self.files[5], self.refs[5]) 

1912 self.assertIsNotNone(cached) 

1913 

1914 # Cache should still report the standard file count. 

1915 self.assertEqual(cache_manager.file_count, threshold + 1) 

1916 

1917 # Add additional entry to cache. 

1918 # Should now have 4, 5, 6 

1919 cached = cache_manager.move_to_cache(self.files[6], self.refs[6]) 

1920 self.assertIsNotNone(cached) 

1921 

1922 # Is the file still there? 

1923 self.assertTrue(found.exists()) 

1924 

1925 # Can we read it? 

1926 data = found.read() 

1927 self.assertGreater(len(data), 0) 

1928 

1929 # Outside context the file should no longer exist. 

1930 self.assertFalse(found.exists()) 

1931 

1932 # File count should not have changed. 

1933 self.assertEqual(cache_manager.file_count, threshold + 1) 

1934 

1935 # Dataset 2 was in the exempt directory but because hardlinks 

1936 # are used it was deleted from the main cache during cache expiry 

1937 # above and so should no longer be found. 

1938 with cache_manager.find_in_cache(self.refs[2], ".txt") as found: 

1939 self.assertIsNone(found) 

1940 

1941 # And the one stored after it is also gone. 

1942 with cache_manager.find_in_cache(self.refs[3], ".txt") as found: 

1943 self.assertIsNone(found) 

1944 

1945 # But dataset 4 is present. 

1946 with cache_manager.find_in_cache(self.refs[4], ".txt") as found: 

1947 self.assertIsNotNone(found) 

1948 

1949 # Adding a new dataset to the cache should now delete it. 

1950 cache_manager.move_to_cache(self.files[7], self.refs[7]) 

1951 

1952 with cache_manager.find_in_cache(self.refs[2], ".txt") as found: 

1953 self.assertIsNone(found) 

1954 

1955 def testCacheExpiryDatasets(self) -> None: 

1956 threshold = 2 # Keep 2 datasets. 

1957 mode = "datasets" 

1958 config_str = self._expiration_config(mode, threshold) 

1959 

1960 cache_manager = self._make_cache_manager(config_str) 

1961 self.assertExpiration(cache_manager, 5, threshold + 1) 

1962 self.assertIn(f"{mode}={threshold}", str(cache_manager)) 

1963 

1964 def testCacheExpiryDatasetsFromDisabled(self) -> None: 

1965 threshold = 2 

1966 mode = "datasets" 

1967 with unittest.mock.patch.dict( 

1968 os.environ, 

1969 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": f"{mode}={threshold}"}, 

1970 ): 

1971 cache_manager = DatastoreCacheManager.create_disabled(universe=DimensionUniverse()) 

1972 self.assertExpiration(cache_manager, 5, threshold + 1) 

1973 self.assertIn(f"{mode}={threshold}", str(cache_manager)) 

1974 

1975 def testExpirationModeOverride(self) -> None: 

1976 threshold = 2 # Keep 2 datasets. 

1977 mode = "datasets" 

1978 config_str = self._expiration_config(mode, threshold) 

1979 

1980 mode = "size" 

1981 threshold = 55 

1982 with unittest.mock.patch.dict( 

1983 os.environ, 

1984 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": f"{mode}={threshold}"}, 

1985 ): 

1986 cache_manager = self._make_cache_manager(config_str) 

1987 self.assertExpiration(cache_manager, 10, 6) 

1988 self.assertIn(f"{mode}={threshold}", str(cache_manager)) 

1989 

1990 # Check we get a warning with unrecognized form. 

1991 with unittest.mock.patch.dict( 

1992 os.environ, 

1993 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": "something"}, 

1994 ): 

1995 with self.assertLogs(level="WARNING") as cm: 

1996 self._make_cache_manager(config_str) 

1997 self.assertIn("Unrecognized form (something)", cm.output[0]) 

1998 

1999 with unittest.mock.patch.dict( 

2000 os.environ, 

2001 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": "something=5"}, 

2002 ): 

2003 with self.assertRaises(ValueError) as cm: 

2004 self._make_cache_manager(config_str) 

2005 self.assertIn("Unrecognized value", str(cm.exception)) 

2006 

2007 def testMissingThreshold(self) -> None: 

2008 threshold = "" 

2009 mode = "datasets" 

2010 config_str = self._expiration_config(mode, threshold) 

2011 

2012 with self.assertRaises(ValueError) as cm: 

2013 self._make_cache_manager(config_str) 

2014 self.assertIn("Cache expiration threshold", str(cm.exception)) 

2015 

2016 def testCacheExpiryDatasetsComposite(self) -> None: 

2017 threshold = 2 # Keep 2 datasets. 

2018 mode = "datasets" 

2019 config_str = self._expiration_config(mode, threshold) 

2020 

2021 cache_manager = self._make_cache_manager(config_str) 

2022 

2023 n_datasets = 3 

2024 for i in range(n_datasets): 

2025 for component_file, component_ref in zip(self.comp_files[i], self.comp_refs[i], strict=True): 

2026 cached = cache_manager.move_to_cache(component_file, component_ref) 

2027 self.assertIsNotNone(cached) 

2028 self.assertTrue(cache_manager.known_to_cache(component_ref)) 

2029 self.assertTrue(cache_manager.known_to_cache(component_ref.makeCompositeRef())) 

2030 self.assertTrue(cache_manager.known_to_cache(component_ref, component_file.getExtension())) 

2031 

2032 self.assertEqual(cache_manager.file_count, 6) # 2 datasets each of 3 files 

2033 

2034 # Write two new non-composite and the number of files should drop. 

2035 self.assertExpiration(cache_manager, 2, 5) 

2036 

2037 def testCacheExpirySize(self) -> None: 

2038 threshold = 55 # Each file is 10 bytes 

2039 mode = "size" 

2040 config_str = self._expiration_config(mode, threshold) 

2041 

2042 cache_manager = self._make_cache_manager(config_str) 

2043 self.assertExpiration(cache_manager, 10, 6) 

2044 self.assertIn(f"{mode}={threshold}", str(cache_manager)) 

2045 

2046 def testDisabledCache(self) -> None: 

2047 # Configure an active cache but disable via environment. 

2048 threshold = 2 

2049 mode = "datasets" 

2050 config_str = self._expiration_config(mode, threshold) 

2051 

2052 with unittest.mock.patch.dict( 

2053 os.environ, 

2054 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": "disabled"}, 

2055 ): 

2056 env_cache_manager = self._make_cache_manager(config_str) 

2057 

2058 # Configure to be disabled 

2059 threshold = 0 

2060 mode = "disabled" 

2061 config_str = self._expiration_config(mode, threshold) 

2062 cfg_cache_manager = self._make_cache_manager(config_str) 

2063 

2064 for cache_manager in ( 

2065 cfg_cache_manager, 

2066 env_cache_manager, 

2067 DatastoreCacheManager.create_disabled(universe=DimensionUniverse()), 

2068 ): 

2069 for uri, ref in zip(self.files, self.refs, strict=True): 

2070 self.assertFalse(cache_manager.should_be_cached(ref)) 

2071 self.assertIsNone(cache_manager.move_to_cache(uri, ref)) 

2072 self.assertFalse(cache_manager.known_to_cache(ref)) 

2073 with cache_manager.find_in_cache(ref, ".txt") as found: 

2074 self.assertIsNone(found, msg=f"{cache_manager}") 

2075 self.assertIn("disabled", str(cache_manager)) 

2076 

2077 def assertExpiration( 

2078 self, cache_manager: DatastoreCacheManager, n_datasets: int, n_retained: int 

2079 ) -> None: 

2080 """Insert the datasets and then check the number retained.""" 

2081 for i in range(n_datasets): 

2082 cached = cache_manager.move_to_cache(self.files[i], self.refs[i]) 

2083 self.assertIsNotNone(cached) 

2084 

2085 self.assertEqual(cache_manager.file_count, n_retained) 

2086 

2087 # The oldest file should not be in the cache any more. 

2088 for i in range(n_datasets): 

2089 with cache_manager.find_in_cache(self.refs[i], ".txt") as found: 

2090 if i >= n_datasets - n_retained: 

2091 self.assertIsInstance(found, ResourcePath) 

2092 else: 

2093 self.assertIsNone(found) 

2094 

2095 def testCacheExpiryAge(self) -> None: 

2096 threshold = 1 # Expire older than 2 seconds 

2097 mode = "age" 

2098 config_str = self._expiration_config(mode, threshold) 

2099 

2100 cache_manager = self._make_cache_manager(config_str) 

2101 self.assertIn(f"{mode}={threshold}", str(cache_manager)) 

2102 

2103 # Insert 3 files, then sleep, then insert more. 

2104 for i in range(2): 

2105 cached = cache_manager.move_to_cache(self.files[i], self.refs[i]) 

2106 self.assertIsNotNone(cached) 

2107 time.sleep(2.0) 

2108 for j in range(4): 

2109 i = 2 + j # Continue the counting 

2110 cached = cache_manager.move_to_cache(self.files[i], self.refs[i]) 

2111 self.assertIsNotNone(cached) 

2112 

2113 # Only the files written after the sleep should exist. 

2114 self.assertEqual(cache_manager.file_count, 4) 

2115 with cache_manager.find_in_cache(self.refs[1], ".txt") as found: 

2116 self.assertIsNone(found) 

2117 with cache_manager.find_in_cache(self.refs[2], ".txt") as found: 

2118 self.assertIsInstance(found, ResourcePath) 

2119 

2120 

2121class NullDatastoreTestCase(DatasetTestHelper, unittest.TestCase): 

2122 """Test the null datastore.""" 

2123 

2124 storageClassFactory = StorageClassFactory() 

2125 

2126 def test_basics(self) -> None: 

2127 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict") 

2128 ref = self.makeDatasetRef("metric", DimensionUniverse().empty, storageClass, {}) 

2129 

2130 null = NullDatastore(None, None) 

2131 

2132 self.assertFalse(null.exists(ref)) 

2133 self.assertFalse(null.knows(ref)) 

2134 knows = null.knows_these([ref]) 

2135 self.assertFalse(knows[ref]) 

2136 null.validateConfiguration(ref) 

2137 

2138 with self.assertRaises(FileNotFoundError): 

2139 null.get(ref) 

2140 with self.assertRaises(NotImplementedError): 

2141 null.put("", ref) 

2142 with self.assertRaises(FileNotFoundError): 

2143 null.getURI(ref) 

2144 with self.assertRaises(FileNotFoundError): 

2145 null.getURIs(ref) 

2146 with self.assertRaises(FileNotFoundError): 

2147 null.getManyURIs([ref]) 

2148 with self.assertRaises(NotImplementedError): 

2149 null.getLookupKeys() 

2150 with self.assertRaises(NotImplementedError): 

2151 null.import_records({}) 

2152 with self.assertRaises(NotImplementedError): 

2153 null.export_records([]) 

2154 with self.assertRaises(NotImplementedError): 

2155 null.export_predicted_records([]) 

2156 with self.assertRaises(NotImplementedError): 

2157 null.export([ref]) 

2158 with self.assertRaises(NotImplementedError): 

2159 null.transfer(null, ref) 

2160 with self.assertRaises(NotImplementedError): 

2161 null.emptyTrash() 

2162 with self.assertRaises(NotImplementedError): 

2163 null.trash(ref) 

2164 with self.assertRaises(NotImplementedError): 

2165 null.forget([ref]) 

2166 with self.assertRaises(NotImplementedError): 

2167 null.remove(ref) 

2168 with self.assertRaises(NotImplementedError): 

2169 null.retrieveArtifacts([ref], ResourcePath(".")) 

2170 with self.assertRaises(NotImplementedError): 

2171 null.transfer_from(null, [ref]) 

2172 with self.assertRaises(NotImplementedError): 

2173 null.ingest() 

2174 

2175 

2176class DatasetRefURIsTestCase(unittest.TestCase): 

2177 """Tests for DatasetRefURIs.""" 

2178 

2179 def testSequenceAccess(self) -> None: 

2180 """Verify that DatasetRefURIs can be treated like a two-item tuple.""" 

2181 uris = DatasetRefURIs() 

2182 

2183 self.assertEqual(len(uris), 2) 

2184 self.assertEqual(uris[0], None) 

2185 self.assertEqual(uris[1], {}) 

2186 

2187 primaryURI = ResourcePath("1/2/3") 

2188 componentURI = ResourcePath("a/b/c") 

2189 

2190 # affirm that DatasetRefURIs does not support MutableSequence functions 

2191 with self.assertRaises(TypeError): 

2192 uris[0] = primaryURI 

2193 with self.assertRaises(TypeError): 

2194 uris[1] = {"foo": componentURI} 

2195 

2196 # but DatasetRefURIs can be set by property name: 

2197 uris.primaryURI = primaryURI 

2198 uris.componentURIs = {"foo": componentURI} 

2199 self.assertEqual(uris.primaryURI, primaryURI) 

2200 self.assertEqual(uris[0], primaryURI) 

2201 

2202 primary, components = uris 

2203 self.assertEqual(primary, primaryURI) 

2204 self.assertEqual(components, {"foo": componentURI}) 

2205 

2206 def testRepr(self) -> None: 

2207 """Verify __repr__ output.""" 

2208 uris = DatasetRefURIs(ResourcePath("/1/2/3"), {"comp": ResourcePath("/a/b/c")}) 

2209 self.assertEqual( 

2210 repr(uris), 

2211 'DatasetRefURIs(ResourcePath("file:///1/2/3"), {\'comp\': ResourcePath("file:///a/b/c")})', 

2212 ) 

2213 

2214 

2215class StoredFileInfoTestCase(DatasetTestHelper, unittest.TestCase): 

2216 """Test the StoredFileInfo class.""" 

2217 

2218 storageClassFactory = StorageClassFactory() 

2219 

2220 def test_StoredFileInfo(self) -> None: 

2221 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict") 

2222 ref = self.makeDatasetRef("metric", DimensionUniverse().empty, storageClass, {}) 

2223 

2224 record = dict( 

2225 storage_class="StructuredDataDict", 

2226 formatter="lsst.daf.butler.Formatter", 

2227 path="a/b/c.txt", 

2228 component="component", 

2229 checksum=None, 

2230 file_size=5, 

2231 ) 

2232 info = StoredFileInfo.from_record(record) 

2233 

2234 self.assertEqual(info.to_record(), record) 

2235 

2236 ref2 = self.makeDatasetRef("metric", DimensionUniverse().empty, storageClass, {}) 

2237 rebased = info.rebase(ref2) 

2238 self.assertEqual(rebased.rebase(ref), info) 

2239 

2240 with self.assertRaises(TypeError): 

2241 rebased.update(formatter=42) 

2242 

2243 with self.assertRaises(ValueError): 

2244 rebased.update(something=42, new="42") 

2245 

2246 # Check that pickle works on StoredFileInfo. 

2247 pickled_info = pickle.dumps(info) 

2248 unpickled_info = pickle.loads(pickled_info) 

2249 self.assertEqual(unpickled_info, info) 

2250 

2251 def test_make_datastore_path_relative(self): 

2252 self.assertEqual(make_datastore_path_relative("a/relative/path"), "a/relative/path") 

2253 self.assertEqual(make_datastore_path_relative("path/with#fragment"), "path/with#fragment") 

2254 self.assertEqual(make_datastore_path_relative("http://server.com/some/path"), "some/path") 

2255 self.assertEqual(make_datastore_path_relative("http://server.com/some/path#frag"), "some/path#frag") 

2256 

2257 def test_datastore_record_data_json_types(self): 

2258 """Test that we don't round-trip checksums to UUIDs when deserializing 

2259 datastore record data. 

2260 """ 

2261 test_json = """ 

2262 { 

2263 "dataset_ids": [ 

2264 "74478304-abf1-4a9c-9eb2-926090a84446" 

2265 ], 

2266 "records": { 

2267 "lsst.daf.butler.datastore.stored_file_info.StoredFileInfo": { 

2268 "74478304abf14a9c9eb2926090a84446": { 

2269 "file_datastore_records": [ 

2270 { 

2271 "formatter": "lsst.daf.butler.formatters.yaml.YamlFormatter", 

2272 "path": "gain_factors/base-2025-158/gain_factors_spx_base-2025-158.yaml", 

2273 "storage_class": "GainFactors", 

2274 "component": "__NULL_STRING__", 

2275 "checksum": "cab515f6-ab67-0484-393f-aaa525dd526f", 

2276 "file_size": 5412 

2277 } 

2278 ] 

2279 } 

2280 } 

2281 } 

2282 } 

2283 """ 

2284 id_str = "74478304abf14a9c9eb2926090a84446" 

2285 s = SerializedDatastoreRecordData.model_validate_json(test_json) 

2286 self.assertIsInstance( 

2287 s.records[get_full_type_name(StoredFileInfo)][id_str]["file_datastore_records"][0]["checksum"], 

2288 str, 

2289 ) 

2290 id = uuid.UUID(id_str) 

2291 d = DatastoreRecordData.from_simple(s) 

2292 self.assertIsInstance(d.records[id]["file_datastore_records"][0], StoredFileInfo) 

2293 self.assertIsInstance(d.records[id]["file_datastore_records"][0].checksum, str) 

2294 

2295 

2296@contextlib.contextmanager 

2297def _temp_yaml_file(data: Any) -> Iterator[str]: 

2298 fh = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml") 

2299 try: 

2300 yaml.dump(data, stream=fh) 

2301 fh.flush() 

2302 yield fh.name 

2303 finally: 

2304 # Some tests delete the file 

2305 with contextlib.suppress(FileNotFoundError): 

2306 fh.close() 

2307 

2308 

2309if __name__ == "__main__": 

2310 unittest.main()