Coverage for tests/test_predicted_qg.py: 13%

215 statements  

« prev     ^ index     » next       coverage.py v7.11.0, created at 2025-11-25 00:45 -0800

1# This file is part of pipe_base. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://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 <https://www.gnu.org/licenses/>. 

27 

28import itertools 

29import logging 

30import os 

31import shutil 

32import subprocess 

33import tempfile 

34import unittest 

35import uuid 

36from io import StringIO 

37 

38import lsst.utils.tests 

39from lsst.pipe.base import QuantumGraph 

40from lsst.pipe.base.dot_tools import graph2dot 

41from lsst.pipe.base.mermaid_tools import graph2mermaid 

42from lsst.pipe.base.quantum_graph import ( 

43 FORMAT_VERSION, 

44 AddressReader, 

45 PredictedDatasetInfo, 

46 PredictedQuantumGraph, 

47 PredictedQuantumInfo, 

48) 

49from lsst.pipe.base.tests.mocks import DynamicConnectionConfig, InMemoryRepo 

50 

51logging.getLogger("timer").setLevel(logging.INFO) 

52logging.getLogger("lsst").setLevel(logging.INFO) 

53logging.getLogger("lsst.pipe.base.quantum_graph").setLevel(logging.DEBUG) 

54 

55 

56class PredictedQuantumGraphTestCase(unittest.TestCase): 

57 """Unit tests for the predicted quantum graph classes. 

58 

59 These tests focus on the interface of the predicted quantum graph 

60 especially I/O and conversion to and from the old quantum graph class. 

61 

62 Additional test coverage is provided by the tests for quantum graph 

63 builders and executors, which all use this class as well. 

64 """ 

65 

66 def setUp(self): 

67 self.helper = InMemoryRepo("base.yaml", "spatial.yaml") 

68 self.helper.add_task( 

69 "calibrate", 

70 dimensions=["visit", "detector"], 

71 inputs={ 

72 "input_image": DynamicConnectionConfig( 

73 dataset_type_name="raw", 

74 dimensions=["visit", "detector"], 

75 ) 

76 }, 

77 prerequisite_inputs={ 

78 "refcat": DynamicConnectionConfig( 

79 dataset_type_name="references", 

80 dimensions=["htm7"], 

81 multiple=True, 

82 ) 

83 }, 

84 init_outputs={ 

85 "output_schema": DynamicConnectionConfig( 

86 dataset_type_name="source_schema", 

87 ) 

88 }, 

89 outputs={ 

90 "output_image": DynamicConnectionConfig( 

91 dataset_type_name="image", 

92 dimensions=["visit", "detector"], 

93 ), 

94 "output_table": DynamicConnectionConfig( 

95 dataset_type_name="source_detector", 

96 dimensions=["visit", "detector"], 

97 ), 

98 }, 

99 ) 

100 self.helper.add_task( 

101 "consolidate", 

102 dimensions=["visit"], 

103 init_inputs={ 

104 "input_schema": DynamicConnectionConfig( 

105 dataset_type_name="source_schema", 

106 ) 

107 }, 

108 inputs={ 

109 "input_table": DynamicConnectionConfig( 

110 dataset_type_name="source_detector", 

111 dimensions=["visit", "detector"], 

112 multiple=True, 

113 ) 

114 }, 

115 outputs={ 

116 "output_table": DynamicConnectionConfig( 

117 dataset_type_name="source", 

118 dimensions=["visit"], 

119 ) 

120 }, 

121 ) 

122 self.helper.add_task( 

123 "resample", 

124 dimensions=["patch", "visit"], 

125 inputs={ 

126 "input_image": DynamicConnectionConfig( 

127 dataset_type_name="image", 

128 dimensions=["visit", "detector"], 

129 multiple=True, 

130 ) 

131 }, 

132 outputs={ 

133 "output_image": DynamicConnectionConfig( 

134 dataset_type_name="warp", 

135 dimensions=["patch", "visit"], 

136 ) 

137 }, 

138 ) 

139 self.helper.add_task( 

140 "coadd", 

141 dimensions=["patch", "band"], 

142 inputs={ 

143 "input_image": DynamicConnectionConfig( 

144 dataset_type_name="warp", 

145 dimensions=["patch", "visit"], 

146 multiple=True, 

147 ) 

148 }, 

149 outputs={ 

150 "output_image": DynamicConnectionConfig( 

151 dataset_type_name="coadd", 

152 dimensions=["patch", "band"], 

153 ), 

154 "output_table": DynamicConnectionConfig( 

155 dataset_type_name="object", 

156 dimensions=["patch", "band"], 

157 # Like all other (defaulted) storage classes here, 

158 # 'ArrowAstropy' below will be mocked; we pick it so we can 

159 # try a component input, and because we know its pytype is 

160 # safe to import here. 

161 storage_class="ArrowAstropy", 

162 ), 

163 }, 

164 ) 

165 self.helper.add_task( 

166 "match", 

167 dimensions=["htm6"], 

168 inputs={ 

169 "input_object": DynamicConnectionConfig( 

170 dataset_type_name="object", 

171 dimensions=["patch", "band"], 

172 multiple=True, 

173 storage_class="ArrowAstropy", 

174 ), 

175 "input_source": DynamicConnectionConfig( 

176 dataset_type_name="source", 

177 dimensions=["visit"], 

178 multiple=True, 

179 ), 

180 "input_object_schema": DynamicConnectionConfig( 

181 dataset_type_name="object.schema", 

182 dimensions=["visit"], 

183 multiple=True, 

184 storage_class="ArrowAstropySchema", 

185 ), 

186 }, 

187 outputs={ 

188 "output_table": DynamicConnectionConfig( 

189 dataset_type_name="matches", 

190 dimensions=["htm6"], 

191 ) 

192 }, 

193 ) 

194 self.builder = self.helper.make_quantum_graph_builder() 

195 

196 def check_quantum_graph( 

197 self, 

198 qg: PredictedQuantumGraph, 

199 *, 

200 all_tasks: bool = True, 

201 thin_graph: bool = True, 

202 all_quantum_datasets: bool = True, 

203 init_quanta: bool = True, 

204 dimension_data_loaded: bool = True, 

205 dimension_data_deserialized: bool = True, 

206 converting_partial: bool = False, 

207 ) -> None: 

208 """Test the attributes and methods of a quantum graph produced by this 

209 test case's builder. 

210 

211 Parameters 

212 ---------- 

213 qg : `lsst.pipe.base.quantum_graph.PredictedQuantumGraph` 

214 Graph to test. 

215 all_tasks : `bool`, optional 

216 Whether all tasks in the pipeline were loaded. 

217 thin_graph : `bool`, optional 

218 Whether the ``thin_graph`` component was loaded. 

219 all_quantum_datasets : `bool`, optional 

220 Whether all quantum datasets were loaded. 

221 init_quanta : `bool`, optional 

222 Whether the ``init_quanta`` component was loaded. 

223 dimension_data_loaded : `bool`, optional 

224 Whether the ``dimension_data`` component was loaded at all. 

225 dimension_data_deserialized : `bool`, optional 

226 Whether all dimension records are expected to have been 

227 deserialized. Ignored if ``dimension_data_loaded`` is `False`. 

228 converting_partial : `bool`, optional 

229 Whether this was a partial load that only included some quanta. 

230 """ 

231 self.assertEqual(qg.header.inputs, ["input_chain"]) 

232 self.assertEqual(qg.header.output, "output_chain") 

233 self.assertEqual(qg.header.output_run, "output_run") 

234 self.assertEqual(qg.header.metadata["stuff"], "whatever") 

235 self.assertEqual(qg.header.version, FORMAT_VERSION) 

236 if not converting_partial: 

237 # When partial-reading an old QG, counts reflect what was loaded. 

238 # In all other cases they reflect the full original graph. 

239 self.assertEqual(qg.header.n_task_quanta["calibrate"], 8) 

240 self.assertEqual(qg.header.n_task_quanta["consolidate"], 2) 

241 self.assertEqual(qg.header.n_task_quanta["resample"], 10) 

242 self.assertEqual(qg.header.n_task_quanta["coadd"], 10) 

243 self.assertEqual(qg.header.n_task_quanta["match"], 2) 

244 self.assertEqual(qg.header.n_quanta, 32) 

245 self.assertEqual( 

246 qg.header.n_datasets, 

247 63 # non-automatic datasets (see breakdown below) 

248 + 32 * 2 # one metadata and one log for each quantam 

249 + 5 # one config for each task 

250 + 1, # packages 

251 ) 

252 if all_tasks: 

253 self.assertFalse(self.helper.pipeline_graph.diff_tasks(qg.pipeline_graph)) 

254 if thin_graph or all_quantum_datasets: 

255 self.assertEqual(len(qg), qg.header.n_quanta) 

256 self.assertEqual(len(qg.quantum_only_xgraph), qg.header.n_quanta) 

257 self.assertEqual(len(qg.quanta_by_task["calibrate"]), 8) 

258 self.assertEqual(len(qg.quanta_by_task["consolidate"]), 2) 

259 self.assertEqual(len(qg.quanta_by_task["resample"]), 10) 

260 self.assertEqual(len(qg.quanta_by_task["coadd"]), 10) 

261 self.assertEqual(len(qg.quanta_by_task["match"]), 2) 

262 if init_quanta: 

263 if all_tasks: 

264 self.assertEqual(len(qg.datasets_by_type["calibrate_config"]), 1) 

265 self.assertEqual(len(qg.datasets_by_type["consolidate_config"]), 1) 

266 self.assertEqual(len(qg.datasets_by_type["resample_config"]), 1) 

267 self.assertEqual(len(qg.datasets_by_type["coadd_config"]), 1) 

268 self.assertEqual(len(qg.datasets_by_type["match_config"]), 1) 

269 self.assertEqual(len(qg.datasets_by_type["source_schema"]), 1) 

270 self.assertEqual(len(qg.datasets_by_type["packages"]), 1) 

271 n_init_datasets = 7 

272 else: 

273 n_init_datasets = 0 

274 if all_quantum_datasets: 

275 assert all_tasks 

276 self.assertEqual( 

277 len(qg.bipartite_xgraph), 

278 qg.header.n_quanta + qg.header.n_datasets - n_init_datasets, 

279 ) 

280 self.assertEqual(len(qg.datasets_by_type["raw"]), 8) 

281 self.assertEqual(len(qg.datasets_by_type["references"]), 4) 

282 self.assertEqual(len(qg.datasets_by_type["image"]), 8) 

283 self.assertEqual(len(qg.datasets_by_type["source_detector"]), 8) 

284 self.assertEqual(len(qg.datasets_by_type["source"]), 2) 

285 self.assertEqual(len(qg.datasets_by_type["warp"]), 10) 

286 self.assertEqual(len(qg.datasets_by_type["coadd"]), 10) 

287 self.assertEqual(len(qg.datasets_by_type["object"]), 10) 

288 self.assertEqual(len(qg.datasets_by_type["matches"]), 2) 

289 # Spot-check some edges and graph structure. 

290 for data_id, dataset_id in qg.datasets_by_type["source_detector"].items(): 

291 for producer_id, _, edge_data in qg.bipartite_xgraph.in_edges(dataset_id, data=True): 

292 self.assertIn(producer_id, qg.quanta_by_task["calibrate"].values()) 

293 self.assertFalse(edge_data["is_read"]) 

294 self.assertEqual( 

295 edge_data["pipeline_edges"], 

296 [qg.pipeline_graph.tasks["calibrate"].get_output_edge("output_table")], 

297 ) 

298 for _, consumer_id, edge_data in qg.bipartite_xgraph.out_edges(dataset_id, data=True): 

299 self.assertIn(consumer_id, qg.quanta_by_task["consolidate"].values()) 

300 self.assertTrue(edge_data["is_read"]) 

301 self.assertEqual( 

302 edge_data["pipeline_edges"], 

303 [qg.pipeline_graph.tasks["consolidate"].get_input_edge("input_table")], 

304 ) 

305 # We use 'is' checks instead of just equality because we don't want 

306 # duplicates of these objects floating around wasting memory. 

307 for task_label, quanta_for_task in qg.quanta_by_task.items(): 

308 for data_id, quantum_id in quanta_for_task.items(): 

309 d1: PredictedQuantumInfo = qg.quantum_only_xgraph.nodes[quantum_id] 

310 self.assertEqual(d1["task_label"], task_label) 

311 self.assertIs(d1["data_id"], data_id) 

312 self.assertIs(d1["pipeline_node"], qg.pipeline_graph.tasks[task_label]) 

313 d2: PredictedQuantumInfo = qg.bipartite_xgraph.nodes[quantum_id] 

314 self.assertEqual(d2["task_label"], task_label) 

315 self.assertIs(d2["data_id"], data_id) 

316 self.assertIs(d2["pipeline_node"], qg.pipeline_graph.tasks[task_label]) 

317 for dataset_type_name, datasets_for_type in qg.datasets_by_type.items(): 

318 if ( 

319 dataset_type_name == "source_schema" 

320 or dataset_type_name == "packages" 

321 or dataset_type_name.endswith("_config") 

322 ): 

323 continue 

324 for data_id, dataset_id in datasets_for_type.items(): 

325 d3: PredictedDatasetInfo = qg.bipartite_xgraph.nodes[dataset_id] 

326 self.assertIs(d3["data_id"], data_id) 

327 self.assertIs(d3["pipeline_node"], qg.pipeline_graph.dataset_types[dataset_type_name]) 

328 if dimension_data_loaded: 

329 self.assertIsNotNone(qg.dimension_data) 

330 if dimension_data_deserialized and not converting_partial: 

331 self.assertEqual(len(qg.dimension_data.records["instrument"]), 1) 

332 self.assertEqual(len(qg.dimension_data.records["visit"]), 2) 

333 self.assertEqual(len(qg.dimension_data.records["detector"]), 4) 

334 self.assertEqual(len(qg.dimension_data.records["visit_detector_region"]), 8) 

335 self.assertEqual(len(qg.dimension_data.records["tract"]), 2) 

336 self.assertEqual(len(qg.dimension_data.records["patch"]), 9) 

337 self.assertEqual(len(qg.dimension_data.records["band"]), 2) 

338 self.assertEqual(len(qg.dimension_data.records["physical_filter"]), 2) 

339 

340 def define_partial_read(self, qg: PredictedQuantumGraph) -> list[uuid.UUID]: 

341 """Return a list of quantum IDs to test partial loads. 

342 

343 Parameters 

344 ---------- 

345 qg : `lsst.pipe.base.quantum_graph.PredictedQuantumGraph` 

346 Full graph to obtain quantum IDs from. Note that quantum IDs are 

347 random and change every test run, and hence test behavior should 

348 never depend on explicit values or their sort order. 

349 

350 Returns 

351 ------- 

352 quanta : `list` [ `uuid.UUID` ] 

353 Quanta to use for partial-read tests. These always have the same 

354 task labels and data IDs: all ``calibrate`` and ``consolidate`` 

355 quanta for ``visit==1``. 

356 """ 

357 return [ 

358 quantum_id 

359 for data_id, quantum_id in qg.quanta_by_task["calibrate"].items() 

360 if data_id["visit"] == 1 

361 ] + [ 

362 quantum_id 

363 for data_id, quantum_id in qg.quanta_by_task["consolidate"].items() 

364 if data_id["visit"] == 1 

365 ] 

366 

367 def check_partial_read( 

368 self, qg: PredictedQuantumGraph, quanta: list[uuid.UUID], converting: bool = False 

369 ) -> None: 

370 """Test the attributes and methods of a partially-loaded quantum graph. 

371 

372 Parameters 

373 ---------- 

374 qg : `lsst.pipe.base.quantum_graph.PredictedQuantumGraph` 

375 Graph for which only ``quanta`` were loaded. 

376 quanta : `list` [ `uuid.UUID` ] 

377 IDs of the quanta that were loaded. Must be the result of a call 

378 to `define_partial_read`. 

379 converting : `bool`, optional 

380 Whether this graph was read via the old 

381 `lsst.pipe.base.QuantumGraph` class and then converted. 

382 """ 

383 self.check_quantum_graph( 

384 qg, 

385 all_tasks=False, 

386 thin_graph=False, 

387 all_quantum_datasets=False, 

388 converting_partial=converting, 

389 dimension_data_deserialized=converting, 

390 ) 

391 self.assertEqual(len(qg), 5) 

392 self.assertEqual(len(qg.quantum_only_xgraph), 5) 

393 execution_quanta = qg.build_execution_quanta() 

394 self.assertEqual(execution_quanta.keys(), set(quanta)) 

395 for quantum_id, quantum in execution_quanta.items(): 

396 self.assertIs(qg.bipartite_xgraph.nodes[quantum_id]["quantum"], quantum) 

397 self.assertIs(qg.quantum_only_xgraph.nodes[quantum_id]["quantum"], quantum) 

398 

399 def test_build(self) -> None: 

400 """Test building a `PredictedQuantumGraph` by 

401 inspecting the result. 

402 """ 

403 qg = self.builder.finish( 

404 output="output_chain", 

405 metadata={"stuff": "whatever"}, 

406 attach_datastore_records=False, 

407 ).assemble() 

408 self.check_quantum_graph(qg) 

409 

410 def test_from_old_quantum_graph(self) -> None: 

411 """Test building a old `QuantumGraph` and then 

412 converting it to a `PredictedQuantumGraph`. 

413 """ 

414 old_qg = self.builder.build( 

415 metadata={ 

416 "input": ["input_chain"], 

417 "output": "output_chain", 

418 "output_run": "output_run", 

419 "stuff": "whatever", 

420 }, 

421 attach_datastore_records=False, 

422 ) 

423 new_qg = PredictedQuantumGraph.from_old_quantum_graph(old_qg) 

424 self.check_quantum_graph(new_qg) 

425 

426 def test_read_execution_quanta_old_file(self) -> None: 

427 """Test building a old `QuantumGraph`, saving it, and then reading it 

428 via `PredictedQuantumGraph.read_execution_quanta`. 

429 """ 

430 old_qg = self.builder.build( 

431 metadata={ 

432 "input": ["input_chain"], 

433 "output": "output_chain", 

434 "output_run": "output_run", 

435 "stuff": "whatever", 

436 }, 

437 attach_datastore_records=False, 

438 ) 

439 with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir: 

440 tmpfile = os.path.join(tmpdir, "old.qgraph") 

441 old_qg.saveUri(tmpfile) 

442 # Read everything. 

443 qg1 = PredictedQuantumGraph.read_execution_quanta(tmpfile) 

444 self.check_quantum_graph(qg1, thin_graph=False) 

445 # Read just a few quanta. 

446 quanta = self.define_partial_read(qg1) 

447 qg2 = PredictedQuantumGraph.read_execution_quanta(tmpfile, quantum_ids=quanta) 

448 self.check_partial_read(qg2, quanta, converting=True) 

449 

450 def test_roundtrip_old_quantum_graph(self) -> None: 

451 """Test building a `PredictedQuantumGraph` and round-tripping it 

452 through the old `QuantumGraph` class. 

453 """ 

454 qg1 = self.builder.finish( 

455 output="output_chain", 

456 metadata={"stuff": "whatever"}, 

457 attach_datastore_records=False, 

458 ).assemble() 

459 old_qg = qg1.to_old_quantum_graph() 

460 qg2 = PredictedQuantumGraph.from_old_quantum_graph(old_qg) 

461 self.check_quantum_graph(qg2) 

462 

463 def test_write_new_as_old(self) -> None: 

464 """Test building a `PredictedQuantumGraphComponents`, saving it with 

465 the old extension, and reading it back in both the old class and the 

466 new class. 

467 """ 

468 components = self.builder.finish( 

469 output="output_chain", 

470 metadata={"stuff": "whatever"}, 

471 attach_datastore_records=False, 

472 ) 

473 with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir: 

474 tmpfile = os.path.join(tmpdir, "old.qgraph") 

475 components.write(tmpfile) 

476 qg1 = PredictedQuantumGraph.read_execution_quanta(tmpfile) 

477 self.check_quantum_graph(qg1) 

478 old_qg = QuantumGraph.loadUri(tmpfile) 

479 qg2 = PredictedQuantumGraph.from_old_quantum_graph(old_qg) 

480 self.check_quantum_graph(qg2) 

481 

482 def test_read_new_as_old(self) -> None: 

483 """Test building a `PredictedQuantumGraphComponents`, saving it with 

484 the new extension, and reading it back in via the old class. 

485 """ 

486 components = self.builder.finish( 

487 output="output_chain", 

488 metadata={"stuff": "whatever"}, 

489 attach_datastore_records=False, 

490 ) 

491 with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir: 

492 tmpfile = os.path.join(tmpdir, "new.qg") 

493 components.write(tmpfile) 

494 old_qg = QuantumGraph.loadUri(tmpfile) 

495 new_qg = PredictedQuantumGraph.from_old_quantum_graph(old_qg) 

496 self.check_quantum_graph(new_qg) 

497 

498 def test_io(self) -> None: 

499 """Test saving a `PredictedQuantumGraphComponents` and reading it back 

500 in as a `PredictedQuantumGraph`, both fully and partially, and as a 

501 `QuantumGraph`, both fully and partially. 

502 """ 

503 components = self.builder.finish( 

504 output="output_chain", 

505 metadata={"stuff": "whatever"}, 

506 attach_datastore_records=False, 

507 ) 

508 # We use a small page size since this is a tiny graph and hence 

509 # otherwise all addresses would fit in a single page. 

510 four_row_page_size = AddressReader.compute_row_size(int_size=8, n_addresses=1) * 4 

511 with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir: 

512 tmpfile = os.path.join(tmpdir, "new.qg") 

513 components.write(tmpfile, zstd_dict_n_inputs=24) # enable dict compression code path 

514 # Test a full read with the new class. 

515 with PredictedQuantumGraph.open(tmpfile, page_size=four_row_page_size) as reader: 

516 reader.read_all() 

517 full_qg = reader.finish() 

518 self.check_quantum_graph(full_qg, dimension_data_deserialized=False) 

519 # Test a full read with the old class (uses new class and then 

520 # converts to old, and we convert back to new for the test). 

521 self.check_quantum_graph( 

522 PredictedQuantumGraph.from_old_quantum_graph(QuantumGraph.loadUri(tmpfile)) 

523 ) 

524 # Test a partial read with the old class. 

525 quanta = self.define_partial_read(full_qg) 

526 self.check_partial_read( 

527 PredictedQuantumGraph.from_old_quantum_graph(QuantumGraph.loadUri(tmpfile, nodes=quanta)), 

528 quanta, 

529 converting=True, 

530 ) 

531 # Test a thin but shallow read with the new class. 

532 with PredictedQuantumGraph.open(tmpfile, page_size=four_row_page_size) as reader: 

533 reader.read_thin_graph() 

534 thin_qg = reader.finish() 

535 self.check_quantum_graph( 

536 thin_qg, 

537 dimension_data_deserialized=False, 

538 all_quantum_datasets=False, 

539 dimension_data_loaded=False, 

540 init_quanta=False, 

541 ) 

542 # Test a deep read of just a few quanta with the new class. 

543 narrow_qg = PredictedQuantumGraph.read_execution_quanta( 

544 tmpfile, 

545 quantum_ids=quanta, 

546 page_size=four_row_page_size, 

547 ) 

548 self.check_partial_read(narrow_qg, quanta) 

549 

550 def test_no_compression_dict(self) -> None: 

551 """Test saving a `PredictedQuantumGraphComponents` and reading it back 

552 in as a `PredictedQuantumGraph` with the compression dictionary 

553 disabled. 

554 """ 

555 components = self.builder.finish( 

556 output="output_chain", 

557 metadata={"stuff": "whatever"}, 

558 attach_datastore_records=False, 

559 ) 

560 # We use a small page size since this is a tiny graph and hence 

561 # otherwise all addresses would fit in a single page. 

562 three_row_page_size = AddressReader.compute_row_size(int_size=8, n_addresses=1) * 3 

563 with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir: 

564 tmpfile = os.path.join(tmpdir, "new.qg") 

565 components.write(tmpfile, zstd_dict_size=0) 

566 with PredictedQuantumGraph.open(tmpfile, page_size=three_row_page_size) as reader: 

567 reader.read_all() 

568 full_qg = reader.finish() 

569 self.check_quantum_graph(full_qg, dimension_data_deserialized=False) 

570 

571 def test_dot(self) -> None: 

572 """Test visualization via GraphViz dot.""" 

573 qg = self.builder.finish(attach_datastore_records=False).assemble() 

574 stream = StringIO() 

575 graph2dot(qg, stream) 

576 if (dot_cmd := shutil.which("dot")) is None: 

577 raise unittest.SkipTest("Aborting test early; 'dot' command is not available.") 

578 # Just check that the dot command can parse what we've given it. 

579 result = subprocess.run( 

580 [dot_cmd, "-Txdot"], input=stream.getvalue().encode(), stdout=subprocess.DEVNULL 

581 ) 

582 result.check_returncode() 

583 

584 def test_mermaid(self) -> None: 

585 """Test visualization via Mermaid.""" 

586 qg = self.builder.finish(attach_datastore_records=False).assemble() 

587 stream = StringIO() 

588 # Just check that it runs without error. 

589 graph2mermaid(qg, stream) 

590 

591 def test_update_output_run(self) -> None: 

592 """Test the PredictedQuantumGraphComponents.output_run method.""" 

593 components = self.builder.finish( 

594 output="output_chain", 

595 metadata={"stuff": "whatever"}, 

596 attach_datastore_records=False, 

597 ) 

598 components.update_output_run("new_output_run") 

599 qg = components.assemble() 

600 for ref in qg.get_init_outputs("calibrate").values(): 

601 self.assertEqual(ref.run, "new_output_run") 

602 for quantum in qg.build_execution_quanta().values(): 

603 for ref in itertools.chain.from_iterable(quantum.outputs.values()): 

604 self.assertEqual(ref.run, "new_output_run", msg=str(ref)) 

605 for ref in itertools.chain.from_iterable(quantum.inputs.values()): 

606 if ref.datasetType.name not in ("raw", "references"): 

607 self.assertEqual(ref.run, "new_output_run", msg=str(ref)) 

608 

609 

610if __name__ == "__main__": 

611 lsst.utils.tests.init() 

612 unittest.main()