Coverage for tests/test_dimension_record_containers.py: 11%

300 statements  

« prev     ^ index     » next       coverage.py v7.11.0, created at 2025-11-04 01:36 -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 

28import copy 

29import json 

30import unittest 

31 

32import pyarrow as pa 

33import pyarrow.parquet as pq 

34import pydantic 

35 

36from lsst.daf.butler import ( 

37 DataCoordinate, 

38 DimensionDataAttacher, 

39 DimensionDataExtractor, 

40 DimensionRecordSet, 

41 DimensionRecordSetDeserializer, 

42 DimensionRecordTable, 

43 SerializableDimensionData, 

44 SerializedKeyValueDimensionRecord, 

45) 

46from lsst.daf.butler.arrow_utils import TimespanArrowType 

47from lsst.daf.butler.tests.utils import create_populated_sqlite_registry 

48 

49DIMENSION_DATA_FILES = [ 

50 "resource://lsst.daf.butler/tests/registry_data/base.yaml", 

51 "resource://lsst.daf.butler/tests/registry_data/spatial.yaml", 

52] 

53 

54 

55class DimensionRecordContainersTestCase(unittest.TestCase): 

56 """Tests for the DimensionRecordTable class.""" 

57 

58 @classmethod 

59 def setUpClass(cls): 

60 # Create an in-memory SQLite database and Registry just to import the 

61 # YAML data. 

62 cls.butler = create_populated_sqlite_registry(*DIMENSION_DATA_FILES) 

63 cls.records = { 

64 element: tuple(list(cls.butler.registry.queryDimensionRecords(element))) 

65 for element in ("visit", "skymap", "patch") 

66 } 

67 cls.universe = cls.butler.dimensions 

68 cls.data_ids = list(cls.butler.registry.queryDataIds(["visit", "patch"]).expanded()) 

69 

70 def assert_record_sets_equal(self, a: DimensionRecordSet, b: DimensionRecordSet) -> None: 

71 self.assertEqual(a, b) # only checks data ID equality 

72 self.assertCountEqual(list(a), list(b)) 

73 

74 def test_record_table_schema_visit(self): 

75 """Test that the Arrow schema for 'visit' has the right types, 

76 including dictionary encoding. 

77 """ 

78 schema = DimensionRecordTable.make_arrow_schema(self.universe["visit"]) 

79 self.assertEqual(schema.field("instrument").type, pa.dictionary(pa.int32(), pa.string())) 

80 self.assertEqual(schema.field("id").type, pa.uint64()) 

81 self.assertEqual(schema.field("physical_filter").type, pa.dictionary(pa.int32(), pa.string())) 

82 self.assertEqual(schema.field("name").type, pa.string()) 

83 self.assertEqual(schema.field("observation_reason").type, pa.dictionary(pa.int32(), pa.string())) 

84 

85 def test_record_table_schema_skymap(self): 

86 """Test that the Arrow schema for 'skymap' has the right types, 

87 including dictionary encoding. 

88 """ 

89 schema = DimensionRecordTable.make_arrow_schema(self.universe["skymap"]) 

90 self.assertEqual(schema.field("name").type, pa.string()) 

91 self.assertEqual(schema.field("hash").type, pa.binary()) 

92 

93 def test_empty_record_table_visit(self): 

94 """Test methods on a table that was initialized with no records. 

95 

96 We use 'visit' records for this test because they have both timespans 

97 and regions, and those are the tricky column types for interoperability 

98 with Arrow. 

99 """ 

100 table = DimensionRecordTable(self.universe["visit"]) 

101 self.assertEqual(len(table), 0) 

102 self.assertEqual(list(table), []) 

103 self.assertEqual(table.element, self.universe["visit"]) 

104 with self.assertRaises(IndexError): 

105 table[0] 

106 self.assertEqual(len(table.column("instrument")), 0) 

107 self.assertEqual(len(table.column("id")), 0) 

108 self.assertEqual(len(table.column("physical_filter")), 0) 

109 self.assertEqual(len(table.column("name")), 0) 

110 self.assertEqual(len(table.column("day_obs")), 0) 

111 table.extend(self.records["visit"]) 

112 self.assertCountEqual(table, self.records["visit"]) 

113 self.assertEqual( 

114 table.to_arrow().schema, DimensionRecordTable.make_arrow_schema(self.universe["visit"]) 

115 ) 

116 

117 def test_empty_record_table_skymap(self): 

118 """Test methods on a table that was initialized with no records. 

119 

120 We use 'skymap' records for this test because that's the only one with 

121 a "hash" column. 

122 """ 

123 table = DimensionRecordTable(self.universe["skymap"]) 

124 self.assertEqual(len(table), 0) 

125 self.assertEqual(list(table), []) 

126 self.assertEqual(table.element, self.universe["skymap"]) 

127 with self.assertRaises(IndexError): 

128 table[0] 

129 self.assertEqual(len(table.column("name")), 0) 

130 self.assertEqual(len(table.column("hash")), 0) 

131 table.extend(self.records["skymap"]) 

132 self.assertCountEqual(table, self.records["skymap"]) 

133 self.assertEqual( 

134 table.to_arrow().schema, DimensionRecordTable.make_arrow_schema(self.universe["skymap"]) 

135 ) 

136 

137 def test_full_record_table_visit(self): 

138 """Test methods on a table that was initialized with an iterable. 

139 

140 We use 'visit' records for this test because they have both timespans 

141 and regions, and those are the tricky column types for interoperability 

142 with Arrow. 

143 """ 

144 table = DimensionRecordTable(self.universe["visit"], self.records["visit"]) 

145 self.assertEqual(len(table), 2) 

146 self.assertEqual(table[0], self.records["visit"][0]) 

147 self.assertEqual(table[1], self.records["visit"][1]) 

148 self.assertEqual(list(table), list(self.records["visit"])) 

149 self.assertListEqual([x.region for x in table], [x.region for x in self.records["visit"]]) 

150 self.assertListEqual([x.timespan for x in table], [x.timespan for x in self.records["visit"]]) 

151 self.assertEqual(table.element, self.universe["visit"]) 

152 self.assertEqual(table.column("instrument")[0].as_py(), "Cam1") 

153 self.assertEqual(table.column("instrument")[1].as_py(), "Cam1") 

154 self.assertEqual(table.column("id")[0].as_py(), 1) 

155 self.assertEqual(table.column("id")[1].as_py(), 2) 

156 self.assertEqual(table.column("name")[0].as_py(), "1") 

157 self.assertEqual(table.column("name")[1].as_py(), "2") 

158 self.assertEqual(table.column("physical_filter")[0].as_py(), "Cam1-G") 

159 self.assertEqual(table.column("physical_filter")[1].as_py(), "Cam1-R1") 

160 self.assertEqual(table.column("day_obs")[0].as_py(), 20210909) 

161 self.assertEqual(table.column("day_obs")[1].as_py(), 20210909) 

162 self.assertEqual(list(table[:1]), list(self.records["visit"][:1])) 

163 self.assertEqual(list(table[1:]), list(self.records["visit"][1:])) 

164 table.extend(self.records["visit"]) 

165 self.assertCountEqual(table, self.records["visit"] + self.records["visit"]) 

166 

167 def test_full_record_table_skymap(self): 

168 """Test methods on a table that was initialized with an iterable. 

169 

170 We use 'skymap' records for this test because that's the only one with 

171 a "hash" column. 

172 """ 

173 table = DimensionRecordTable(self.universe["skymap"], self.records["skymap"]) 

174 self.assertEqual(len(table), 1) 

175 self.assertEqual(table[0], self.records["skymap"][0]) 

176 self.assertEqual(list(table), list(self.records["skymap"])) 

177 self.assertEqual(table.element, self.universe["skymap"]) 

178 self.assertEqual(table.column("name")[0].as_py(), "SkyMap1") 

179 self.assertEqual(table.column("hash")[0].as_py(), b"notreallyahashofanything!") 

180 table.extend(self.records["skymap"]) 

181 self.assertCountEqual(table, self.records["skymap"] + self.records["skymap"]) 

182 

183 def test_record_table_parquet_visit(self): 

184 """Test round-tripping a dimension record table through Parquet. 

185 

186 We use 'visit' records for this test because they have both timespans 

187 and regions, and those are the tricky column types for interoperability 

188 with Arrow. 

189 """ 

190 original_records = list(self.records["visit"]) 

191 table1 = DimensionRecordTable(self.universe["visit"], self.records["visit"]) 

192 stream = pa.BufferOutputStream() 

193 pq.write_table(table1.to_arrow(), stream) 

194 table2 = DimensionRecordTable( 

195 universe=self.universe, table=pq.read_table(pa.BufferReader(stream.getvalue())) 

196 ) 

197 self.assertEqual(list(table1), list(table2)) 

198 self.assertEqual(original_records, list(table2)) 

199 self.assertListEqual([x.region for x in table1], [x.region for x in table2]) 

200 self.assertListEqual([x.region for x in original_records], [x.region for x in table2]) 

201 self.assertListEqual([x.timespan for x in table1], [x.timespan for x in table2]) 

202 self.assertListEqual([x.timespan for x in original_records], [x.timespan for x in table2]) 

203 

204 def test_record_table_parquet_skymap(self): 

205 """Test round-tripping a dimension record table through Parquet. 

206 

207 We use 'skymap' records for this test because that's the only one with 

208 a "hash" column. 

209 """ 

210 table1 = DimensionRecordTable(self.universe["skymap"], self.records["skymap"]) 

211 stream = pa.BufferOutputStream() 

212 pq.write_table(table1.to_arrow(), stream) 

213 table2 = DimensionRecordTable( 

214 universe=self.universe, table=pq.read_table(pa.BufferReader(stream.getvalue())) 

215 ) 

216 self.assertEqual(list(table1), list(table2)) 

217 

218 def test_record_chunk_init(self): 

219 """Test constructing a DimensionRecordTable from an iterable in chunks. 

220 

221 We use 'patch' records for this test because there are enough of them 

222 to have multiple chunks. 

223 """ 

224 table1 = DimensionRecordTable(self.universe["patch"], self.records["patch"], batch_size=5) 

225 self.assertEqual(len(table1), 12) 

226 self.assertEqual([len(batch) for batch in table1.to_arrow().to_batches()], [5, 5, 2]) 

227 self.assertEqual(list(table1), list(self.records["patch"])) 

228 

229 def test_arrow_timespan_nulls(self): 

230 """Test that null values for Timespan can be round-tripped through 

231 arrow. 

232 """ 

233 schema = pa.schema([pa.field("timespan", TimespanArrowType(), nullable=True)]) 

234 table = pa.Table.from_pylist([{"timespan": None}], schema=schema) 

235 output_value = table.to_pylist()[0]["timespan"] 

236 self.assertIsNone(output_value) 

237 

238 def test_record_set_const(self): 

239 """Test attributes and methods of `DimensionRecordSet` that do not 

240 modify the set. 

241 

242 We use 'patch' records for this test because there are enough of them 

243 to do nontrivial set-operation tests. 

244 """ 

245 element = self.universe["patch"] 

246 records = self.records["patch"] 

247 set1 = DimensionRecordSet(element, records[:7]) 

248 self.assertEqual(set1, DimensionRecordSet("patch", records[:7], universe=self.universe)) 

249 # DimensionRecordSets do not compare as equal with other set types, 

250 # even with the same content. 

251 self.assertNotEqual(set1, set(records[:7])) 

252 with self.assertRaises(TypeError): 

253 DimensionRecordSet("patch", records[:7]) 

254 self.assertEqual(set1.element, self.universe["patch"]) 

255 self.assertEqual(len(set1), 7) 

256 self.assertEqual(list(set1), list(records[:7])) 

257 self.assertIn(records[4], set1) 

258 self.assertIn(records[5].dataId, set1) 

259 self.assertNotIn(self.records["visit"][0], set1) 

260 self.assertTrue(set1.issubset(DimensionRecordSet(element, records[:8]))) 

261 self.assertFalse(set1.issubset(DimensionRecordSet(element, records[1:6]))) 

262 with self.assertRaises(ValueError): 

263 set1.issubset(DimensionRecordSet(self.universe["tract"])) 

264 self.assertTrue(set1.issuperset(DimensionRecordSet(element, records[1:6]))) 

265 self.assertFalse(set1.issuperset(DimensionRecordSet(element, records[:8]))) 

266 with self.assertRaises(ValueError): 

267 set1.issuperset(DimensionRecordSet(self.universe["tract"])) 

268 self.assertTrue(set1.isdisjoint(DimensionRecordSet(element, records[7:]))) 

269 self.assertFalse(set1.isdisjoint(DimensionRecordSet(element, records[5:8]))) 

270 with self.assertRaises(ValueError): 

271 set1.isdisjoint(DimensionRecordSet(self.universe["tract"])) 

272 self.assertEqual( 

273 set1.intersection(DimensionRecordSet(element, records[5:])), 

274 DimensionRecordSet(element, records[5:7]), 

275 ) 

276 self.assertEqual( 

277 set1.intersection(DimensionRecordSet(element, records[5:])), 

278 DimensionRecordSet(element, records[5:7]), 

279 ) 

280 with self.assertRaises(ValueError): 

281 set1.intersection(DimensionRecordSet(self.universe["tract"])) 

282 self.assertEqual( 

283 set1.difference(DimensionRecordSet(element, records[5:])), 

284 DimensionRecordSet(element, records[:5]), 

285 ) 

286 with self.assertRaises(ValueError): 

287 set1.difference(DimensionRecordSet(self.universe["tract"])) 

288 self.assertEqual( 

289 set1.union(DimensionRecordSet(element, records[5:9])), 

290 DimensionRecordSet(element, records[:9]), 

291 ) 

292 with self.assertRaises(ValueError): 

293 set1.union(DimensionRecordSet(self.universe["tract"])) 

294 self.assertEqual(set1.find(records[0].dataId), records[0]) 

295 with self.assertRaises(LookupError): 

296 set1.find(self.records["patch"][8].dataId) 

297 with self.assertRaises(ValueError): 

298 set1.find(self.records["visit"][0].dataId) 

299 self.assertEqual(set1.find_with_required_values(records[0].dataId.required_values), records[0]) 

300 

301 def test_record_set_add(self): 

302 """Test DimensionRecordSet.add.""" 

303 set1 = DimensionRecordSet("patch", self.records["patch"][:2], universe=self.universe) 

304 set1.add(self.records["patch"][2]) 

305 with self.assertRaises(ValueError): 

306 set1.add(self.records["visit"][0]) 

307 self.assertEqual(set1, DimensionRecordSet("patch", self.records["patch"][:3], universe=self.universe)) 

308 set1.add(self.records["patch"][2]) 

309 self.assertEqual(list(set1), list(self.records["patch"][:3])) 

310 

311 def test_record_set_find_or_add(self): 

312 """Test DimensionRecordSet.find and find_with_required_values with 

313 a 'or_add' callback. 

314 """ 

315 set1 = DimensionRecordSet("patch", self.records["patch"][:2], universe=self.universe) 

316 set1.find(self.records["patch"][2].dataId, or_add=lambda _c, _r: self.records["patch"][2]) 

317 with self.assertRaises(ValueError): 

318 set1.find(self.records["visit"][0].dataId, or_add=lambda _c, _r: self.records["visit"][0]) 

319 self.assertEqual(set1, DimensionRecordSet("patch", self.records["patch"][:3], universe=self.universe)) 

320 

321 set1.find_with_required_values( 

322 self.records["patch"][3].dataId.required_values, or_add=lambda _c, _r: self.records["patch"][3] 

323 ) 

324 self.assertEqual(set1, DimensionRecordSet("patch", self.records["patch"][:4], universe=self.universe)) 

325 

326 def test_record_set_update_from_data_coordinates(self): 

327 """Test DimensionRecordSet.update_from_data_coordinates.""" 

328 set1 = DimensionRecordSet("patch", self.records["patch"][:2], universe=self.universe) 

329 set1.update_from_data_coordinates(self.data_ids) 

330 for data_id in self.data_ids: 

331 self.assertIn(data_id.records["patch"], set1) 

332 

333 def test_record_set_discard(self): 

334 """Test DimensionRecordSet.discard.""" 

335 set1 = DimensionRecordSet("patch", self.records["patch"][:2], universe=self.universe) 

336 set2 = copy.deepcopy(set1) 

337 # These discards should do nothing. 

338 set1.discard(self.records["patch"][2]) 

339 self.assertEqual(set1, set2) 

340 set1.discard(self.records["patch"][2].dataId) 

341 self.assertEqual(set1, set2) 

342 with self.assertRaises(ValueError): 

343 set1.discard(self.records["visit"][0]) 

344 self.assertEqual(set1, set2) 

345 with self.assertRaises(ValueError): 

346 set1.discard(self.records["visit"][0].dataId) 

347 self.assertEqual(set1, set2) 

348 # These ones should remove a record from each set. 

349 set1.discard(self.records["patch"][1]) 

350 set2.discard(self.records["patch"][1].dataId) 

351 self.assertEqual(set1, set2) 

352 self.assertNotIn(self.records["patch"][1], set1) 

353 self.assertNotIn(self.records["patch"][1], set2) 

354 

355 def test_record_set_remove(self): 

356 """Test DimensionRecordSet.remove.""" 

357 set1 = DimensionRecordSet("patch", self.records["patch"][:2], universe=self.universe) 

358 set2 = copy.deepcopy(set1) 

359 # These removes should raise with strong exception safety. 

360 with self.assertRaises(KeyError): 

361 set1.remove(self.records["patch"][2]) 

362 self.assertEqual(set1, set2) 

363 with self.assertRaises(KeyError): 

364 set1.remove(self.records["patch"][2].dataId) 

365 self.assertEqual(set1, set2) 

366 with self.assertRaises(ValueError): 

367 set1.remove(self.records["visit"][0]) 

368 self.assertEqual(set1, set2) 

369 with self.assertRaises(ValueError): 

370 set1.remove(self.records["visit"][0].dataId) 

371 self.assertEqual(set1, set2) 

372 # These ones should remove a record from each set. 

373 set1.remove(self.records["patch"][1]) 

374 set2.remove(self.records["patch"][1].dataId) 

375 self.assertEqual(set1, set2) 

376 self.assertNotIn(self.records["patch"][1], set1) 

377 self.assertNotIn(self.records["patch"][1], set2) 

378 

379 def test_record_set_pop(self): 

380 """Test DimensionRecordSet.pop.""" 

381 set1 = DimensionRecordSet("patch", self.records["patch"][:2], universe=self.universe) 

382 set2 = copy.deepcopy(set1) 

383 record1 = set1.pop() 

384 set2.remove(record1) 

385 self.assertNotIn(record1, set1) 

386 self.assertEqual(set1, set2) 

387 record2 = set1.pop() 

388 set2.remove(record2) 

389 self.assertNotIn(record2, set1) 

390 self.assertEqual(set1, set2) 

391 self.assertFalse(set1) 

392 

393 def test_record_key_value_serialization(self): 

394 """Test DimensionRecord.serialize_key_value, deserialize_key, and 

395 deserialize_value. 

396 """ 

397 for element, records in self.records.items(): 

398 with self.subTest(element=element): 

399 for record1 in records: 

400 raw1 = record1.serialize_key_value() 

401 self.assertIsInstance(raw1, list) 

402 raw2 = json.loads(json.dumps(raw1)) 

403 self.assertEqual(raw1, raw2) 

404 key, raw_value = record1.definition.RecordClass.deserialize_key(raw2) 

405 self.assertEqual(key, record1.dataId.required_values) 

406 record2 = record1.definition.RecordClass.deserialize_value(key, raw_value) 

407 self.assertEqual(record1, record2) 

408 

409 def test_record_set_serialization(self): 

410 """Test DimensionRecordSet.serialize_records and deserialize_records, 

411 as well as DimensionRecordSetDeserializer. 

412 """ 

413 adapter = pydantic.TypeAdapter(list[SerializedKeyValueDimensionRecord]) 

414 for element, records in self.records.items(): 

415 with self.subTest(element=element): 

416 rs1 = DimensionRecordSet(element, records, universe=self.universe) 

417 raw = adapter.validate_json(adapter.dump_json(rs1.serialize_records())) 

418 rs2 = DimensionRecordSet(element, universe=self.universe) 

419 rs2.deserialize_records(raw) 

420 self.assert_record_sets_equal(rs1, rs2) 

421 deserializer = DimensionRecordSetDeserializer.from_raw(self.universe[element], raw) 

422 self.assertEqual(len(deserializer), len(records)) 

423 for record1 in rs1: 

424 record2 = deserializer[record1.dataId.required_values] 

425 self.assertEqual(record1, record2) 

426 

427 def test_extract_serialize_attach(self): 

428 """Test serializing expanded data IDs by first extracting their 

429 dimension records, serializing them in normalized form, and then 

430 attaching the records to deserialized data IDs. 

431 """ 

432 dimensions = self.universe.conform(["visit", "patch", "htm5"]) 

433 expanded_data_ids_1: list[DataCoordinate] = [] 

434 pixelization = dimensions.universe.skypix_dimensions["htm5"].pixelization 

435 for data_id in self.data_ids: 

436 for htm_begin, htm_end in pixelization.envelope(data_id.records["patch"].region): 

437 for htm_id in range(htm_begin, htm_end): 

438 new_data_id = self.butler.registry.expandDataId(data_id, htm5=htm_id) 

439 self.assertEqual(new_data_id.dimensions, dimensions) 

440 expanded_data_ids_1.append(new_data_id) 

441 extractor = DimensionDataExtractor.from_dimension_group( 

442 dimensions, 

443 ignore=["tract"], 

444 ignore_cached=True, 

445 include_skypix=False, 

446 ) 

447 extractor.update(expanded_data_ids_1) 

448 thin_data_ids = [ 

449 DataCoordinate.from_full_values(dimensions, data_id.full_values) 

450 for data_id in expanded_data_ids_1 

451 ] 

452 model1 = SerializableDimensionData.from_record_sets(extractor.records.values()) 

453 self.assertEqual(model1.root.keys(), {"visit", "patch", "day_obs"}) 

454 # Delete one patch record from the model to probe more logic branches 

455 # in the deserialization; like the tract records, we'll recover it by 

456 # querying the butler. 

457 del model1.root["patch"][-1] 

458 model2 = SerializableDimensionData.model_validate_json(model1.model_dump_json()) 

459 with self.butler.query() as query: 

460 attacher1 = DimensionDataAttacher( 

461 deserializers=model2.make_deserializers(self.universe), 

462 cache=query._driver._dimension_record_cache, 

463 dimensions=dimensions, 

464 ) 

465 expanded_data_ids_2 = attacher1.attach(dimensions, thin_data_ids, query=query) 

466 # This equality check is necessary but not sufficient for this test, 

467 # because DataCoordinate equality doesn't look at records. 

468 self.assertEqual(expanded_data_ids_1, expanded_data_ids_2) 

469 for a, b in zip(expanded_data_ids_1, expanded_data_ids_2): 

470 self.assertEqual(dict(a.records), dict(b.records)) 

471 # Caching in the attacher should ensure that there is only one copy 

472 # of each dimension record in the new set. 

473 all_values = {(data_id["tract"], data_id["patch"]) for data_id in expanded_data_ids_2} 

474 all_identities = { 

475 id(data_id.records["patch"]): data_id.records["patch"] for data_id in expanded_data_ids_2 

476 } 

477 self.assertGreater(len(all_values), 1) 

478 self.assertEqual(len(all_identities), len(all_values), msg=str(all_identities)) 

479 # Serialize and round-trip again, via attacher.serialized(), and attach 

480 # all records again. 

481 model3 = attacher1.serialized() 

482 attacher2 = DimensionDataAttacher( 

483 deserializers=model3.make_deserializers(self.universe), dimensions=dimensions 

484 ) 

485 expanded_data_ids_4 = attacher2.attach(dimensions, thin_data_ids) 

486 self.assertEqual(expanded_data_ids_1, expanded_data_ids_4) 

487 for a, b in zip(expanded_data_ids_1, expanded_data_ids_4): 

488 self.assertEqual(dict(a.records), dict(b.records)) 

489 # Attach data IDs with the original attacher again, after clearing out 

490 # its deserializers. The cached records should all be present, and it 

491 # should work without any butler queries. 

492 attacher1.deserializers.clear() 

493 expanded_data_ids_3 = attacher1.attach(dimensions, thin_data_ids) 

494 self.assertEqual(expanded_data_ids_1, expanded_data_ids_3) 

495 for a, b in zip(expanded_data_ids_1, expanded_data_ids_3): 

496 self.assertEqual(dict(a.records), dict(b.records)) 

497 # Delete a record and try again to check error-handling. 

498 removed_patch_data_id = next(iter(attacher1.records["patch"])).dataId 

499 attacher1.records["patch"].remove(removed_patch_data_id) 

500 with self.assertRaisesRegex(LookupError, r"No dimension record.*patch.*"): 

501 attacher1.attach(dimensions, thin_data_ids) 

502 

503 

504if __name__ == "__main__": 

505 unittest.main()