Coverage for tests/test_script.py: 99%

1# This file is part of rucio_register 

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 program is free software: you can redistribute it and/or modify 

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

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

12# (at your option) any later version. 

13# 

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

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

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

17# GNU General Public License for more details. 

18# 

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

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

21 

22 

23import json 

24import logging 

25import os 

26import tempfile 

27import unittest 

28from concurrent.futures import Future 

29from unittest.mock import MagicMock, patch 

30 

31from click.testing import CliRunner 

32 

33import lsst.utils.tests 

34from lsst.daf.butler import _exceptions 

35from lsst.rucio.register import script 

36 

37 

38class LimitDatasetTypesTestCase(unittest.TestCase): 

39 def test_limit_dataset_types(self): 

40 dataset_types = ["a", "b", "c", "d"] 

41 

42 self.assertEqual(script._limit_dataset_types(dataset_types, None), dataset_types) 

43 self.assertEqual(script._limit_dataset_types(dataset_types, 2), ["a", "b"]) 

44 self.assertEqual(script._limit_dataset_types(dataset_types, 0), []) 

45 

46 

47class SetLogLevelTestCase(unittest.TestCase): 

48 def test_set_log_level(self): 

49 # String input 

50 script._set_log_level("DEBUG") 

51 log = logging.getLogger() 

52 self.assertEqual(log.level, logging.DEBUG) 

53 self.assertTrue(len(log.handlers) > 0) 

54 

55 # Dict input from daf_butler cli log_level_option 

56 script._set_log_level({None: "INFO"}) 

57 self.assertEqual(log.level, logging.INFO) 

58 

59 

60class ChunksTestCase(unittest.TestCase): 

61 def test_chunks_materialized_list(self): 

62 dataset_refs = list(range(10)) 

63 ref_chunks = list(script.chunks(dataset_refs, 3)) 

64 self.assertEqual(ref_chunks, [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]) 

65 

66 def test_chunks_exact_multiple(self): 

67 dataset_refs = list(range(9)) 

68 ref_chunks = list(script.chunks(dataset_refs, 3)) 

69 self.assertEqual(ref_chunks, [[0, 1, 2], [3, 4, 5], [6, 7, 8]]) 

70 

71 def test_chunks_empty(self): 

72 ref_chunks = list(script.chunks([], 3)) 

73 self.assertEqual(ref_chunks, []) 

74 

75 def test_chunks_chunk_size_larger_than_input(self): 

76 dataset_refs = [1, 2, 3] 

77 ref_chunks = list(script.chunks(dataset_refs, 10)) 

78 self.assertEqual(ref_chunks, [[1, 2, 3]]) 

79 

80 

81class DummyExecutor: 

82 last_instances = [] 

83 

84 def __init__(self, max_workers=None, mp_context=None, initializer=None, initargs=()): 

85 self.mp_context = mp_context 

86 self.initializer = initializer 

87 self.initargs = initargs 

88 DummyExecutor.last_instances.append(self) 

89 if initializer is not None: 89 ↛ exitline 89 didn't return from function '__init__' because the condition on line 89 was always true

90 initializer(*initargs) 

91 

92 def __enter__(self): 

93 return self 

94 

95 def __exit__(self, exc_type, exc_val, exc_tb): 

96 pass 

97 

98 def submit(self, fn, *args, **kwargs): 

99 fut = Future() 

100 try: 

101 res = fn(*args, **kwargs) 

102 fut.set_result(res) 

103 except Exception as e: 

104 fut.set_exception(e) 

105 return fut 

106 

107 

108class AutoRegisterTestCase(lsst.utils.tests.TestCase): 

109 def setUp(self): 

110 self.tmp_dir = tempfile.TemporaryDirectory() 

111 self.out_dir = os.path.join(self.tmp_dir.name, "uuids") 

112 self.repo = "/tmp/test_repo" 

113 

114 def tearDown(self): 

115 script._thread_local.__dict__.pop("butlers", None) 

116 self.tmp_dir.cleanup() 

117 

118 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor) 

119 @patch("lsst.rucio.register.script.Butler") 

120 def test_auto_register_dry_run_no_raw(self, mock_butler_class): 

121 mock_butler = MagicMock() 

122 mock_butler_class.return_value = mock_butler 

123 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"] 

124 

125 mock_info = MagicMock() 

126 mock_info.dataset_types = ["visitSummary", "deepCoadd"] 

127 mock_butler.collections.get_info.return_value = mock_info 

128 

129 ref1 = MagicMock() 

130 ref1.id = "11111111-1111-1111-1111-111111111111" 

131 ref2 = MagicMock() 

132 ref2.id = "22222222-2222-2222-2222-222222222222" 

133 mock_butler.query_datasets.return_value = [ref1, ref2] 

134 

135 runner = CliRunner() 

136 result = runner.invoke( 

137 script.main, 

138 [ 

139 "auto-register", 

140 "--repo", 

141 self.repo, 

142 "--root-chain", 

143 "HSC/runs/*", 

144 "--start-date", 

145 "2023-01-01", 

146 "--cutoff-date", 

147 "2023-04-01", 

148 "--dry-run", 

149 "--out-dir", 

150 self.out_dir, 

151 ], 

152 ) 

153 

154 self.assertEqual(result.exit_code, 0, msg=result.output) 

155 

156 sanitized_repo = self.repo.strip("/").replace("/", "_") 

157 expected_dir = os.path.join(self.out_dir, sanitized_repo, "auto_register") 

158 self.assertTrue(os.path.exists(expected_dir)) 

159 

160 files = os.listdir(expected_dir) 

161 self.assertGreater(len(files), 0) 

162 

163 with open(os.path.join(expected_dir, files[0])) as f: 

164 content = f.read() 

165 

166 self.assertIn("11111111-1111-1111-1111-111111111111", content) 

167 self.assertIn("22222222-2222-2222-2222-222222222222", content) 

168 

169 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor) 

170 @patch("lsst.rucio.register.script._get_rucio_interface") 

171 @patch("lsst.rucio.register.script.Butler") 

172 def test_auto_register_dry_run_with_raw(self, mock_butler_class, mock_get_ri): 

173 mock_butler = MagicMock() 

174 mock_butler_class.return_value = mock_butler 

175 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"] 

176 

177 mock_info = MagicMock() 

178 mock_info.dataset_types = ["raw", "visitSummary"] 

179 mock_butler.collections.get_info.return_value = mock_info 

180 

181 mock_ri = MagicMock() 

182 mock_inner_butler = MagicMock() 

183 mock_get_ri.return_value = (mock_ri, mock_inner_butler) 

184 

185 ref1 = MagicMock() 

186 ref1.id = "33333333-3333-3333-3333-333333333333" 

187 mock_butler.query_datasets.return_value = [ref1] 

188 

189 runner = CliRunner() 

190 result = runner.invoke( 

191 script.main, 

192 [ 

193 "auto-register", 

194 "--repo", 

195 self.repo, 

196 "--root-chain", 

197 "HSC/runs/*", 

198 "--start-date", 

199 "2023-01-01", 

200 "--cutoff-date", 

201 "2023-04-01", 

202 "--dry-run", 

203 "--out-dir", 

204 self.out_dir, 

205 ], 

206 ) 

207 

208 self.assertEqual(result.exit_code, 0, msg=result.output) 

209 sanitized_repo = self.repo.strip("/").replace("/", "_") 

210 expected_dir = os.path.join(self.out_dir, sanitized_repo, "auto_register") 

211 self.assertTrue(os.path.exists(expected_dir)) 

212 files = os.listdir(expected_dir) 

213 self.assertGreater(len(files), 0) 

214 

215 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor) 

216 @patch("lsst.rucio.register.script._register", return_value=2) 

217 @patch("lsst.rucio.register.script.Butler") 

218 def test_auto_register_registration_mode(self, mock_butler_class, mock_register): 

219 mock_butler = MagicMock() 

220 mock_butler_class.return_value = mock_butler 

221 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"] 

222 

223 mock_info = MagicMock() 

224 mock_info.dataset_types = ["visitSummary"] 

225 mock_butler.collections.get_info.return_value = mock_info 

226 

227 ref1 = MagicMock() 

228 ref1.id = "44444444-4444-4444-4444-444444444444" 

229 mock_butler.query_datasets.return_value = [ref1] 

230 

231 runner = CliRunner() 

232 result = runner.invoke( 

233 script.main, 

234 [ 

235 "auto-register", 

236 "--repo", 

237 self.repo, 

238 "--root-chain", 

239 "HSC/runs/*", 

240 "--start-date", 

241 "2023-01-01", 

242 "--cutoff-date", 

243 "2023-04-01", 

244 "--out-dir", 

245 self.out_dir, 

246 ], 

247 ) 

248 

249 self.assertEqual(result.exit_code, 0, msg=result.output) 

250 self.assertTrue(mock_register.called) 

251 

252 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor) 

253 @patch("lsst.rucio.register.script.Butler") 

254 def test_auto_register_empty_collections_and_exceptions(self, mock_butler_class): 

255 mock_butler = MagicMock() 

256 mock_butler_class.return_value = mock_butler 

257 

258 # Empty collections 

259 mock_butler.collections.query.return_value = [] 

260 runner = CliRunner() 

261 result = runner.invoke( 

262 script.main, 

263 [ 

264 "auto-register", 

265 "--repo", 

266 self.repo, 

267 "--root-chain", 

268 "HSC/runs/*", 

269 "--start-date", 

270 "2023-01-01", 

271 "--cutoff-date", 

272 "2023-04-01", 

273 "--dry-run", 

274 "--out-dir", 

275 self.out_dir, 

276 ], 

277 ) 

278 self.assertEqual(result.exit_code, 0, msg=result.output) 

279 

280 # None info or empty dataset_types 

281 mock_butler.collections.query.return_value = ["HSC/runs/empty"] 

282 mock_butler.collections.get_info.return_value = None 

283 result = runner.invoke( 

284 script.main, 

285 [ 

286 "auto-register", 

287 "--repo", 

288 self.repo, 

289 "--root-chain", 

290 "HSC/runs/*", 

291 "--start-date", 

292 "2023-01-01", 

293 "--cutoff-date", 

294 "2023-04-01", 

295 "--dry-run", 

296 "--out-dir", 

297 self.out_dir, 

298 ], 

299 ) 

300 self.assertEqual(result.exit_code, 0, msg=result.output) 

301 

302 # mock_info with dataset_types = None 

303 mock_info = MagicMock() 

304 mock_info.dataset_types = None 

305 mock_butler.collections.get_info.return_value = mock_info 

306 result = runner.invoke( 

307 script.main, 

308 [ 

309 "auto-register", 

310 "--repo", 

311 self.repo, 

312 "--root-chain", 

313 "HSC/runs/*", 

314 "--start-date", 

315 "2023-01-01", 

316 "--cutoff-date", 

317 "2023-04-01", 

318 "--dry-run", 

319 "--out-dir", 

320 self.out_dir, 

321 ], 

322 ) 

323 self.assertEqual(result.exit_code, 0, msg=result.output) 

324 

325 # EmptyQueryResultError 

326 mock_info = MagicMock() 

327 mock_info.dataset_types = ["visitSummary"] 

328 mock_butler.collections.get_info.return_value = mock_info 

329 mock_butler.query_datasets.side_effect = _exceptions.EmptyQueryResultError( 

330 reasons=["No datasets found"] 

331 ) 

332 result = runner.invoke( 

333 script.main, 

334 [ 

335 "auto-register", 

336 "--repo", 

337 self.repo, 

338 "--root-chain", 

339 "HSC/runs/*", 

340 "--start-date", 

341 "2023-01-01", 

342 "--cutoff-date", 

343 "2023-04-01", 

344 "--dry-run", 

345 "--out-dir", 

346 self.out_dir, 

347 ], 

348 ) 

349 self.assertEqual(result.exit_code, 0, msg=result.output) 

350 

351 # Generic Exception 

352 mock_butler.query_datasets.side_effect = Exception("Query error") 

353 result = runner.invoke( 

354 script.main, 

355 [ 

356 "auto-register", 

357 "--repo", 

358 self.repo, 

359 "--root-chain", 

360 "HSC/runs/*", 

361 "--start-date", 

362 "2023-01-01", 

363 "--cutoff-date", 

364 "2023-04-01", 

365 "--dry-run", 

366 "--out-dir", 

367 self.out_dir, 

368 ], 

369 ) 

370 self.assertEqual(result.exit_code, 0, msg=result.output) 

371 

372 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor) 

373 @patch("lsst.rucio.register.script._get_rucio_interface") 

374 @patch("lsst.rucio.register.script.Butler") 

375 def test_auto_register_max_dataset_types(self, mock_butler_class, mock_get_ri): 

376 mock_butler = MagicMock() 

377 mock_butler_class.return_value = mock_butler 

378 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"] 

379 

380 mock_info = MagicMock() 

381 mock_info.dataset_types = ["typeA", "typeB", "typeC"] 

382 mock_butler.collections.get_info.return_value = mock_info 

383 

384 mock_ri = MagicMock() 

385 mock_inner_butler = MagicMock() 

386 mock_get_ri.return_value = (mock_ri, mock_inner_butler) 

387 mock_butler.query_datasets.return_value = [] 

388 

389 runner = CliRunner() 

390 result = runner.invoke( 

391 script.main, 

392 [ 

393 "auto-register", 

394 "--repo", 

395 self.repo, 

396 "--root-chain", 

397 "HSC/runs/*", 

398 "--start-date", 

399 "2023-01-01", 

400 "--cutoff-date", 

401 "2023-04-01", 

402 "--max-dataset-types", 

403 "1", 

404 "--dry-run", 

405 "--out-dir", 

406 self.out_dir, 

407 ], 

408 ) 

409 

410 self.assertEqual(result.exit_code, 0, msg=result.output) 

411 called_dataset_types = [ 

412 call.kwargs.get("dataset_type") for call in mock_butler.query_datasets.call_args_list 

413 ] 

414 self.assertIn("typeA", called_dataset_types) 

415 self.assertNotIn("typeB", called_dataset_types) 

416 self.assertNotIn("typeC", called_dataset_types) 

417 

418 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor) 

419 @patch("lsst.rucio.register.script.Butler") 

420 def test_auto_register_direct_call(self, mock_butler_class): 

421 mock_butler = MagicMock() 

422 mock_butler_class.return_value = mock_butler 

423 mock_butler.collections.query.return_value = [] 

424 

425 script.auto_register.callback( 

426 repo=self.repo, 

427 root_chain="HSC/runs/*", 

428 start_date="2023-01-01", 

429 cutoff_date="2023-04-01", 

430 df_name="USDF", 

431 max_did_per_dataset=50000, 

432 dry_run=True, 

433 out_dir=self.out_dir, 

434 log_level={}, 

435 ) 

436 mock_butler.collections.query.assert_called_once() 

437 

438 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor) 

439 @patch("lsst.rucio.register.script._get_rucio_interface") 

440 @patch("lsst.rucio.register.script.Butler") 

441 def test_auto_register_with_transfer_list_file(self, mock_butler_class, mock_get_ri): 

442 yaml_content = """ 

443 stage1: 

444 step1a: 

445 - typeA 

446 - typeC 

447 """ 

448 temp_yaml_path = os.path.join(self.tmp_dir.name, "test_transfer_list.yaml") 

449 with open(temp_yaml_path, "w") as f: 

450 f.write(yaml_content) 

451 

452 mock_butler = MagicMock() 

453 mock_butler_class.return_value = mock_butler 

454 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"] 

455 

456 mock_info = MagicMock() 

457 mock_info.dataset_types = ["raw", "typeA", "typeB", "typeC"] 

458 mock_butler.collections.get_info.return_value = mock_info 

459 

460 mock_ri = MagicMock() 

461 mock_inner_butler = MagicMock() 

462 mock_get_ri.return_value = (mock_ri, mock_inner_butler) 

463 mock_butler.query_datasets.return_value = [] 

464 

465 runner = CliRunner() 

466 result = runner.invoke( 

467 script.main, 

468 [ 

469 "auto-register", 

470 "--repo", 

471 self.repo, 

472 "--root-chain", 

473 "HSC/runs/*", 

474 "--start-date", 

475 "2023-01-01", 

476 "--cutoff-date", 

477 "2023-04-01", 

478 "--transfer-list", 

479 temp_yaml_path, 

480 "--dry-run", 

481 "--out-dir", 

482 self.out_dir, 

483 ], 

484 ) 

485 

486 self.assertEqual(result.exit_code, 0, msg=result.output) 

487 called_dataset_types = [ 

488 call.kwargs.get("dataset_type") for call in mock_butler.query_datasets.call_args_list 

489 ] 

490 self.assertIn("typeA", called_dataset_types) 

491 self.assertIn("typeC", called_dataset_types) 

492 self.assertNotIn("typeB", called_dataset_types) 

493 self.assertNotIn("raw", called_dataset_types) 

494 

495 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor) 

496 @patch("lsst.rucio.register.script.urllib.request.urlopen") 

497 @patch("lsst.rucio.register.script._get_rucio_interface") 

498 @patch("lsst.rucio.register.script.Butler") 

499 def test_auto_register_with_transfer_list_url(self, mock_butler_class, mock_get_ri, mock_urlopen): 

500 mock_response = MagicMock() 

501 mock_response.read.return_value = b""" 

502 stage1: 

503 step1b: 

504 - typeB 

505 """ 

506 mock_urlopen.return_value.__enter__.return_value = mock_response 

507 

508 mock_butler = MagicMock() 

509 mock_butler_class.return_value = mock_butler 

510 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"] 

511 

512 mock_info = MagicMock() 

513 mock_info.dataset_types = ["typeA", "typeB", "typeC"] 

514 mock_butler.collections.get_info.return_value = mock_info 

515 

516 mock_ri = MagicMock() 

517 mock_inner_butler = MagicMock() 

518 mock_get_ri.return_value = (mock_ri, mock_inner_butler) 

519 

520 refA = MagicMock() 

521 refA.id = "11111111-1111-1111-1111-111111111111" 

522 refA.dataset_type.name = "typeA" 

523 refB = MagicMock() 

524 refB.id = "22222222-2222-2222-2222-222222222222" 

525 refB.dataset_type.name = "typeB" 

526 refC = MagicMock() 

527 refC.id = "33333333-3333-3333-3333-333333333333" 

528 refC.dataset_type.name = "typeC" 

529 

530 mock_butler.query_datasets.side_effect = lambda **kw: [ 

531 r for r in (refA, refB, refC) if r.dataset_type.name == kw.get("dataset_type") 

532 ] 

533 

534 runner = CliRunner() 

535 result = runner.invoke( 

536 script.main, 

537 [ 

538 "auto-register", 

539 "--repo", 

540 self.repo, 

541 "--root-chain", 

542 "HSC/runs/*", 

543 "--start-date", 

544 "2023-01-01", 

545 "--cutoff-date", 

546 "2023-04-01", 

547 "--transfer-list", 

548 "https://raw.githubusercontent.com/test/cm_transfer_list.yaml", 

549 "--dry-run", 

550 "--out-dir", 

551 self.out_dir, 

552 ], 

553 ) 

554 

555 self.assertEqual(result.exit_code, 0, msg=result.output) 

556 sanitized_repo = self.repo.strip("/").replace("/", "_") 

557 expected_dir = os.path.join(self.out_dir, sanitized_repo, "auto_register") 

558 self.assertTrue(os.path.exists(expected_dir)) 

559 files = os.listdir(expected_dir) 

560 self.assertGreater(len(files), 0) 

561 

562 with open(os.path.join(expected_dir, files[0])) as f: 

563 content = f.read() 

564 

565 self.assertIn("22222222-2222-2222-2222-222222222222", content) 

566 self.assertNotIn("11111111-1111-1111-1111-111111111111", content) 

567 self.assertNotIn("33333333-3333-3333-3333-333333333333", content) 

568 

569 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor) 

570 @patch("lsst.rucio.register.script._get_rucio_interface") 

571 @patch("lsst.rucio.register.script.Butler") 

572 def test_auto_register_with_dataset_name_prefix(self, mock_butler_class, mock_get_ri): 

573 mock_butler = MagicMock() 

574 mock_butler_class.return_value = mock_butler 

575 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"] 

576 

577 mock_info = MagicMock() 

578 mock_info.dataset_types = ["typeA"] 

579 mock_butler.collections.get_info.return_value = mock_info 

580 

581 mock_ri = MagicMock() 

582 mock_inner_butler = MagicMock() 

583 mock_get_ri.return_value = (mock_ri, mock_inner_butler) 

584 

585 refA = MagicMock() 

586 refA.id = "11111111-1111-1111-1111-111111111111" 

587 refA.dataset_type.name = "typeA" 

588 

589 mock_butler.query_datasets.return_value = [refA] 

590 

591 runner = CliRunner() 

592 result = runner.invoke( 

593 script.main, 

594 [ 

595 "auto-register", 

596 "--repo", 

597 self.repo, 

598 "--root-chain", 

599 "HSC/runs/*", 

600 "--start-date", 

601 "2023-01-01", 

602 "--cutoff-date", 

603 "2023-04-01", 

604 "--dataset-name-prefix", 

605 "Dataset/HSC/runs/RC2/w_2026_23/DM-55175/LANCS/outputs", 

606 "--dry-run", 

607 "--out-dir", 

608 self.out_dir, 

609 ], 

610 ) 

611 

612 self.assertEqual(result.exit_code, 0, msg=result.output) 

613 sanitized_repo = self.repo.strip("/").replace("/", "_") 

614 expected_dir = os.path.join(self.out_dir, sanitized_repo, "auto_register") 

615 self.assertTrue(os.path.exists(expected_dir)) 

616 files = os.listdir(expected_dir) 

617 self.assertTrue( 

618 any("Dataset-HSC-runs-RC2-w_2026_23-DM-55175-LANCS-outputs-typeA-USDF" in f for f in files) 

619 ) 

620 

621 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor) 

622 @patch("lsst.rucio.register.script._process_auto_register_batch") 

623 @patch("lsst.rucio.register.script.Butler") 

624 def test_auto_register_with_clear_is_new(self, mock_butler_class, mock_process_batch): 

625 mock_butler = MagicMock() 

626 mock_butler_class.return_value = mock_butler 

627 mock_process_batch.return_value = {"registered": 1, "failed": 0} 

628 

629 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"] 

630 mock_info = MagicMock() 

631 mock_info.dataset_types = ["typeA"] 

632 mock_butler.collections.get_info.return_value = mock_info 

633 

634 refA = MagicMock() 

635 refA.id = "11111111-1111-1111-1111-111111111111" 

636 refA.dataset_type.name = "typeA" 

637 mock_butler.query_datasets.return_value = [refA] 

638 

639 runner = CliRunner() 

640 result = runner.invoke( 

641 script.main, 

642 [ 

643 "auto-register", 

644 "--repo", 

645 self.repo, 

646 "--root-chain", 

647 "HSC/runs/*", 

648 "--start-date", 

649 "2023-01-01", 

650 "--cutoff-date", 

651 "2023-04-01", 

652 "--clear-is-new", 

653 ], 

654 ) 

655 

656 self.assertEqual(result.exit_code, 0, msg=result.output) 

657 self.assertTrue(mock_process_batch.called) 

658 _, kwargs = mock_process_batch.call_args 

659 self.assertTrue(kwargs.get("clear_is_new")) 

660 

661 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor) 

662 @patch("lsst.rucio.register.script.Butler") 

663 def test_auto_register_mid_quarter_dates(self, mock_butler_class): 

664 mock_butler = MagicMock() 

665 mock_butler_class.return_value = mock_butler 

666 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"] 

667 

668 mock_info = MagicMock() 

669 mock_info.dataset_types = ["typeA"] 

670 mock_butler.collections.get_info.return_value = mock_info 

671 mock_butler.query_datasets.return_value = [] 

672 

673 runner = CliRunner() 

674 result = runner.invoke( 

675 script.main, 

676 [ 

677 "auto-register", 

678 "--repo", 

679 self.repo, 

680 "--root-chain", 

681 "HSC/runs/*", 

682 "--start-date", 

683 "2026-05-15", 

684 "--cutoff-date", 

685 "2026-08-15", 

686 "--dry-run", 

687 "--out-dir", 

688 self.out_dir, 

689 ], 

690 ) 

691 

692 self.assertEqual(result.exit_code, 0, msg=result.output) 

693 where_clauses = [call.kwargs.get("where") for call in mock_butler.query_datasets.call_args_list] 

694 self.assertEqual( 

695 where_clauses, 

696 [ 

697 "ingest_date >= T'2026-05-15T00:00:00' AND ingest_date < T'2026-07-01T00:00:00'", 

698 "ingest_date >= T'2026-07-01T00:00:00' AND ingest_date < T'2026-08-15T00:00:00'", 

699 ], 

700 ) 

701 

702 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor) 

703 @patch("lsst.rucio.register.script._process_auto_register_batch") 

704 @patch("lsst.rucio.register.script.Butler") 

705 def test_auto_register_passes_backoff_parameters(self, mock_butler_class, mock_process_batch): 

706 mock_butler = MagicMock() 

707 mock_butler_class.return_value = mock_butler 

708 mock_process_batch.return_value = {"registered": 1, "failed": 0} 

709 

710 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"] 

711 mock_info = MagicMock() 

712 mock_info.dataset_types = ["typeA"] 

713 mock_butler.collections.get_info.return_value = mock_info 

714 

715 refA = MagicMock() 

716 refA.id = "11111111-1111-1111-1111-111111111111" 

717 refA.dataset_type.name = "typeA" 

718 mock_butler.query_datasets.return_value = [refA] 

719 

720 runner = CliRunner() 

721 result = runner.invoke( 

722 script.main, 

723 [ 

724 "auto-register", 

725 "--repo", 

726 self.repo, 

727 "--root-chain", 

728 "HSC/runs/*", 

729 "--start-date", 

730 "2023-01-01", 

731 "--cutoff-date", 

732 "2023-04-01", 

733 "--backoff-factor", 

734 "2.5", 

735 "--backoff-max-value", 

736 "45", 

737 "--backoff-max-tries", 

738 "12", 

739 ], 

740 ) 

741 

742 self.assertEqual(result.exit_code, 0, msg=result.output) 

743 self.assertTrue(mock_process_batch.called) 

744 _, kwargs = mock_process_batch.call_args 

745 self.assertEqual(kwargs.get("backoff_factor"), 2.5) 

746 self.assertEqual(kwargs.get("backoff_max_value"), 45) 

747 self.assertEqual(kwargs.get("max_retries"), 12) 

748 

749 @patch("lsst.rucio.register.script._get_rucio_interface") 

750 def test_dataset_list_with_json_uuidlist(self, mock_get_ri): 

751 mock_ri = MagicMock() 

752 mock_butler = MagicMock() 

753 mock_get_ri.return_value = (mock_ri, mock_butler) 

754 

755 json_file = os.path.join(self.tmp_dir.name, "auto-register-failures.json") 

756 with open(json_file, "w") as f: 

757 json.dump(["019ee079-b7e4-700d-a0bf-03412ae4fe0a"], f) 

758 

759 mock_ref = MagicMock() 

760 mock_ref.id = "019ee079-b7e4-700d-a0bf-03412ae4fe0a" 

761 mock_butler.get_many_datasets.return_value = [mock_ref] 

762 

763 runner = CliRunner() 

764 result = runner.invoke( 

765 script.main, 

766 [ 

767 "dataset-list", 

768 "--repo", 

769 self.repo, 

770 "--rucio-dataset", 

771 "Dataset/Test/Failed", 

772 "--uuidlist", 

773 json_file, 

774 ], 

775 ) 

776 

777 self.assertEqual(result.exit_code, 0, msg=result.output) 

778 mock_butler.get_many_datasets.assert_called_once_with(["019ee079-b7e4-700d-a0bf-03412ae4fe0a"]) 

779 

780 @patch("lsst.rucio.register.script.RucioInterface") 

781 @patch("lsst.rucio.register.script.Butler") 

782 @patch("lsst.rucio.register.script.RucioRegisterConfig") 

783 def test_get_rucio_interface_timeout(self, mock_config_class, mock_butler_class, mock_ri_class): 

784 script._rucio_interface_cache.clear() 

785 mock_ri_class.DEFAULT_BACKOFF_FACTOR = 5.0 

786 mock_ri_class.DEFAULT_BACKOFF_MAX_VALUE = 30 

787 mock_ri_class.DEFAULT_RETRIES = 8 

788 mock_ri_class.DEFAULT_TIMEOUT = 300.0 

789 

790 mock_config = MagicMock() 

791 mock_config.rucio_rse = "RSE" 

792 mock_config.scope = "test" 

793 mock_config.rse_root = "/tmp" 

794 mock_config.dtn_url = "root://localhost" 

795 mock_config_class.return_value = mock_config 

796 

797 script._get_rucio_interface(self.repo, "dummy.yaml", "DATA_PRODUCT", kwargs={"timeout": 450.0}) 

798 self.assertEqual(mock_ri_class.call_args.kwargs.get("timeout"), 450.0) 

799 script._rucio_interface_cache.clear() 

800 

801 @patch("lsst.rucio.register.script.RucioInterface") 

802 @patch("lsst.rucio.register.script.Butler") 

803 @patch("lsst.rucio.register.script.RucioRegisterConfig") 

804 def test_get_rucio_interface_cache_isolated_by_pid( 

805 self, mock_config_class, mock_butler_class, mock_ri_class 

806 ): 

807 script._rucio_interface_cache.clear() 

808 mock_config = MagicMock() 

809 mock_config.rucio_rse = "RSE" 

810 mock_config.scope = "test" 

811 mock_config.rse_root = "/tmp" 

812 mock_config.dtn_url = "root://localhost" 

813 mock_config_class.return_value = mock_config 

814 

815 mock_butler_class.side_effect = lambda repo: MagicMock(name=f"Butler-{repo}") 

816 mock_ri_class.side_effect = lambda **kw: MagicMock(name="RucioInterface") 

817 

818 with patch("os.getpid", return_value=1000):