Coverage for tests/test_script.py: 99%
317 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-11-03 01:33 -0800
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/>.
23import json
24import logging
25import os
26import tempfile
27import unittest
28from concurrent.futures import Future
29from unittest.mock import MagicMock, patch
31from click.testing import CliRunner
33import lsst.utils.tests
34from lsst.daf.butler import _exceptions
35from lsst.rucio.register import script
38class LimitDatasetTypesTestCase(unittest.TestCase):
39 def test_limit_dataset_types(self):
40 dataset_types = ["a", "b", "c", "d"]
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), [])
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)
55 # Dict input from daf_butler cli log_level_option
56 script._set_log_level({None: "INFO"})
57 self.assertEqual(log.level, logging.INFO)
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]])
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]])
71 def test_chunks_empty(self):
72 ref_chunks = list(script.chunks([], 3))
73 self.assertEqual(ref_chunks, [])
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]])
81class DummyExecutor:
82 def __init__(self, max_workers=None):
83 pass
85 def __enter__(self):
86 return self
88 def __exit__(self, exc_type, exc_val, exc_tb):
89 pass
91 def submit(self, fn, *args, **kwargs):
92 fut = Future()
93 try:
94 res = fn(*args, **kwargs)
95 fut.set_result(res)
96 except Exception as e:
97 fut.set_exception(e)
98 return fut
101class AutoRegisterTestCase(lsst.utils.tests.TestCase):
102 def setUp(self):
103 self.tmp_dir = tempfile.TemporaryDirectory()
104 self.out_dir = os.path.join(self.tmp_dir.name, "uuids")
105 self.repo = "/tmp/test_repo"
107 def tearDown(self):
108 self.tmp_dir.cleanup()
110 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor)
111 @patch("lsst.rucio.register.script.Butler")
112 def test_auto_register_dry_run_no_raw(self, mock_butler_class):
113 mock_butler = MagicMock()
114 mock_butler_class.return_value = mock_butler
115 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"]
117 mock_info = MagicMock()
118 mock_info.dataset_types = ["visitSummary", "deepCoadd"]
119 mock_butler.collections.get_info.return_value = mock_info
121 ref1 = MagicMock()
122 ref1.id = "11111111-1111-1111-1111-111111111111"
123 ref2 = MagicMock()
124 ref2.id = "22222222-2222-2222-2222-222222222222"
125 mock_butler.query_all_datasets.return_value = [ref1, ref2]
127 runner = CliRunner()
128 result = runner.invoke(
129 script.main,
130 [
131 "auto-register",
132 "--repo",
133 self.repo,
134 "--root-chain",
135 "HSC/runs/*",
136 "--start-date",
137 "2023-01-01",
138 "--cutoff-date",
139 "2023-04-01",
140 "--dry-run",
141 "--out-dir",
142 self.out_dir,
143 ],
144 )
146 self.assertEqual(result.exit_code, 0, msg=result.output)
148 sanitized_repo = self.repo.strip("/").replace("/", "_")
149 expected_dir = os.path.join(self.out_dir, sanitized_repo, "auto_register")
150 self.assertTrue(os.path.exists(expected_dir))
152 files = os.listdir(expected_dir)
153 self.assertGreater(len(files), 0)
155 with open(os.path.join(expected_dir, files[0])) as f:
156 content = f.read()
158 self.assertIn("11111111-1111-1111-1111-111111111111", content)
159 self.assertIn("22222222-2222-2222-2222-222222222222", content)
161 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor)
162 @patch("lsst.rucio.register.script._get_rucio_interface")
163 @patch("lsst.rucio.register.script.Butler")
164 def test_auto_register_dry_run_with_raw(self, mock_butler_class, mock_get_ri):
165 mock_butler = MagicMock()
166 mock_butler_class.return_value = mock_butler
167 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"]
169 mock_info = MagicMock()
170 mock_info.dataset_types = ["raw", "visitSummary"]
171 mock_butler.collections.get_info.return_value = mock_info
173 mock_ri = MagicMock()
174 mock_inner_butler = MagicMock()
175 mock_get_ri.return_value = (mock_ri, mock_inner_butler)
177 ref1 = MagicMock()
178 ref1.id = "33333333-3333-3333-3333-333333333333"
179 mock_butler.query_datasets.return_value = [ref1]
181 runner = CliRunner()
182 result = runner.invoke(
183 script.main,
184 [
185 "auto-register",
186 "--repo",
187 self.repo,
188 "--root-chain",
189 "HSC/runs/*",
190 "--start-date",
191 "2023-01-01",
192 "--cutoff-date",
193 "2023-04-01",
194 "--dry-run",
195 "--out-dir",
196 self.out_dir,
197 ],
198 )
200 self.assertEqual(result.exit_code, 0, msg=result.output)
201 sanitized_repo = self.repo.strip("/").replace("/", "_")
202 expected_dir = os.path.join(self.out_dir, sanitized_repo, "auto_register")
203 self.assertTrue(os.path.exists(expected_dir))
204 files = os.listdir(expected_dir)
205 self.assertGreater(len(files), 0)
207 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor)
208 @patch("lsst.rucio.register.script._register", return_value=2)
209 @patch("lsst.rucio.register.script.Butler")
210 def test_auto_register_registration_mode(self, mock_butler_class, mock_register):
211 mock_butler = MagicMock()
212 mock_butler_class.return_value = mock_butler
213 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"]
215 mock_info = MagicMock()
216 mock_info.dataset_types = ["visitSummary"]
217 mock_butler.collections.get_info.return_value = mock_info
219 ref1 = MagicMock()
220 ref1.id = "44444444-4444-4444-4444-444444444444"
221 mock_butler.query_all_datasets.return_value = [ref1]
223 runner = CliRunner()
224 result = runner.invoke(
225 script.main,
226 [
227 "auto-register",
228 "--repo",
229 self.repo,
230 "--root-chain",
231 "HSC/runs/*",
232 "--start-date",
233 "2023-01-01",
234 "--cutoff-date",
235 "2023-04-01",
236 "--out-dir",
237 self.out_dir,
238 ],
239 )
241 self.assertEqual(result.exit_code, 0, msg=result.output)
242 self.assertTrue(mock_register.called)
244 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor)
245 @patch("lsst.rucio.register.script.Butler")
246 def test_auto_register_empty_collections_and_exceptions(self, mock_butler_class):
247 mock_butler = MagicMock()
248 mock_butler_class.return_value = mock_butler
250 # Empty collections
251 mock_butler.collections.query.return_value = []
252 runner = CliRunner()
253 result = runner.invoke(
254 script.main,
255 [
256 "auto-register",
257 "--repo",
258 self.repo,
259 "--root-chain",
260 "HSC/runs/*",
261 "--start-date",
262 "2023-01-01",
263 "--cutoff-date",
264 "2023-04-01",
265 "--dry-run",
266 "--out-dir",
267 self.out_dir,
268 ],
269 )
270 self.assertEqual(result.exit_code, 0, msg=result.output)
272 # None info or empty dataset_types
273 mock_butler.collections.query.return_value = ["HSC/runs/empty"]
274 mock_butler.collections.get_info.return_value = None
275 result = runner.invoke(
276 script.main,
277 [
278 "auto-register",
279 "--repo",
280 self.repo,
281 "--root-chain",
282 "HSC/runs/*",
283 "--start-date",
284 "2023-01-01",
285 "--cutoff-date",
286 "2023-04-01",
287 "--dry-run",
288 "--out-dir",
289 self.out_dir,
290 ],
291 )
292 self.assertEqual(result.exit_code, 0, msg=result.output)
294 # mock_info with dataset_types = None
295 mock_info = MagicMock()
296 mock_info.dataset_types = None
297 mock_butler.collections.get_info.return_value = mock_info
298 result = runner.invoke(
299 script.main,
300 [
301 "auto-register",
302 "--repo",
303 self.repo,
304 "--root-chain",
305 "HSC/runs/*",
306 "--start-date",
307 "2023-01-01",
308 "--cutoff-date",
309 "2023-04-01",
310 "--dry-run",
311 "--out-dir",
312 self.out_dir,
313 ],
314 )
315 self.assertEqual(result.exit_code, 0, msg=result.output)
317 # EmptyQueryResultError
318 mock_info = MagicMock()
319 mock_info.dataset_types = ["visitSummary"]
320 mock_butler.collections.get_info.return_value = mock_info
321 mock_butler.query_all_datasets.side_effect = _exceptions.EmptyQueryResultError(
322 reasons=["No datasets found"]
323 )
324 result = runner.invoke(
325 script.main,
326 [
327 "auto-register",
328 "--repo",
329 self.repo,
330 "--root-chain",
331 "HSC/runs/*",
332 "--start-date",
333 "2023-01-01",
334 "--cutoff-date",
335 "2023-04-01",
336 "--dry-run",
337 "--out-dir",
338 self.out_dir,
339 ],
340 )
341 self.assertEqual(result.exit_code, 0, msg=result.output)
343 # Generic Exception
344 mock_butler.query_all_datasets.side_effect = Exception("Query error")
345 result = runner.invoke(
346 script.main,
347 [
348 "auto-register",
349 "--repo",
350 self.repo,
351 "--root-chain",
352 "HSC/runs/*",
353 "--start-date",
354 "2023-01-01",
355 "--cutoff-date",
356 "2023-04-01",
357 "--dry-run",
358 "--out-dir",
359 self.out_dir,
360 ],
361 )
362 self.assertEqual(result.exit_code, 0, msg=result.output)
364 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor)
365 @patch("lsst.rucio.register.script._get_rucio_interface")
366 @patch("lsst.rucio.register.script.Butler")
367 def test_auto_register_max_dataset_types(self, mock_butler_class, mock_get_ri):
368 mock_butler = MagicMock()
369 mock_butler_class.return_value = mock_butler
370 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"]
372 mock_info = MagicMock()
373 mock_info.dataset_types = ["typeA", "typeB", "typeC"]
374 mock_butler.collections.get_info.return_value = mock_info
376 mock_ri = MagicMock()
377 mock_inner_butler = MagicMock()
378 mock_get_ri.return_value = (mock_ri, mock_inner_butler)
379 mock_butler.query_datasets.return_value = []
381 runner = CliRunner()
382 result = runner.invoke(
383 script.main,
384 [
385 "auto-register",
386 "--repo",
387 self.repo,
388 "--root-chain",
389 "HSC/runs/*",
390 "--start-date",
391 "2023-01-01",
392 "--cutoff-date",
393 "2023-04-01",
394 "--max-dataset-types",
395 "1",
396 "--dry-run",
397 "--out-dir",
398 self.out_dir,
399 ],
400 )
402 self.assertEqual(result.exit_code, 0, msg=result.output)
403 called_dataset_types = [
404 call.kwargs.get("dataset_type") for call in mock_butler.query_datasets.call_args_list
405 ]
406 self.assertIn("typeA", called_dataset_types)
407 self.assertNotIn("typeB", called_dataset_types)
408 self.assertNotIn("typeC", called_dataset_types)
410 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor)
411 @patch("lsst.rucio.register.script.Butler")
412 def test_auto_register_direct_call(self, mock_butler_class):
413 mock_butler = MagicMock()
414 mock_butler_class.return_value = mock_butler
415 mock_butler.collections.query.return_value = []
417 script.auto_register.callback(
418 repo=self.repo,
419 root_chain="HSC/runs/*",
420 start_date="2023-01-01",
421 cutoff_date="2023-04-01",
422 df_name="USDF",
423 max_did_per_dataset=50000,
424 dry_run=True,
425 out_dir=self.out_dir,
426 log_level={},
427 )
428 mock_butler.collections.query.assert_called_once()
430 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor)
431 @patch("lsst.rucio.register.script._get_rucio_interface")
432 @patch("lsst.rucio.register.script.Butler")
433 def test_auto_register_with_transfer_list_file(self, mock_butler_class, mock_get_ri):
434 yaml_content = """
435 stage1:
436 step1a:
437 - typeA
438 - typeC
439 """
440 temp_yaml_path = os.path.join(self.tmp_dir.name, "test_transfer_list.yaml")
441 with open(temp_yaml_path, "w") as f:
442 f.write(yaml_content)
444 mock_butler = MagicMock()
445 mock_butler_class.return_value = mock_butler
446 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"]
448 mock_info = MagicMock()
449 mock_info.dataset_types = ["raw", "typeA", "typeB", "typeC"]
450 mock_butler.collections.get_info.return_value = mock_info
452 mock_ri = MagicMock()
453 mock_inner_butler = MagicMock()
454 mock_get_ri.return_value = (mock_ri, mock_inner_butler)
455 mock_butler.query_datasets.return_value = []
457 runner = CliRunner()
458 result = runner.invoke(
459 script.main,
460 [
461 "auto-register",
462 "--repo",
463 self.repo,
464 "--root-chain",
465 "HSC/runs/*",
466 "--start-date",
467 "2023-01-01",
468 "--cutoff-date",
469 "2023-04-01",
470 "--transfer-list",
471 temp_yaml_path,
472 "--dry-run",
473 "--out-dir",
474 self.out_dir,
475 ],
476 )
478 self.assertEqual(result.exit_code, 0, msg=result.output)
479 called_dataset_types = [
480 call.kwargs.get("dataset_type") for call in mock_butler.query_datasets.call_args_list
481 ]
482 self.assertIn("typeA", called_dataset_types)
483 self.assertIn("typeC", called_dataset_types)
484 self.assertNotIn("typeB", called_dataset_types)
485 self.assertNotIn("raw", called_dataset_types)
487 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor)
488 @patch("lsst.rucio.register.script.urllib.request.urlopen")
489 @patch("lsst.rucio.register.script._get_rucio_interface")
490 @patch("lsst.rucio.register.script.Butler")
491 def test_auto_register_with_transfer_list_url(self, mock_butler_class, mock_get_ri, mock_urlopen):
492 mock_response = MagicMock()
493 mock_response.read.return_value = b"""
494 stage1:
495 step1b:
496 - typeB
497 """
498 mock_urlopen.return_value.__enter__.return_value = mock_response
500 mock_butler = MagicMock()
501 mock_butler_class.return_value = mock_butler
502 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"]
504 mock_info = MagicMock()
505 mock_info.dataset_types = ["typeA", "typeB", "typeC"]
506 mock_butler.collections.get_info.return_value = mock_info
508 mock_ri = MagicMock()
509 mock_inner_butler = MagicMock()
510 mock_get_ri.return_value = (mock_ri, mock_inner_butler)
512 refA = MagicMock()
513 refA.id = "11111111-1111-1111-1111-111111111111"
514 refA.dataset_type.name = "typeA"
515 refB = MagicMock()
516 refB.id = "22222222-2222-2222-2222-222222222222"
517 refB.dataset_type.name = "typeB"
518 refC = MagicMock()
519 refC.id = "33333333-3333-3333-3333-333333333333"
520 refC.dataset_type.name = "typeC"
522 mock_butler.query_all_datasets.return_value = [refA, refB, refC]
524 runner = CliRunner()
525 result = runner.invoke(
526 script.main,
527 [
528 "auto-register",
529 "--repo",
530 self.repo,
531 "--root-chain",
532 "HSC/runs/*",
533 "--start-date",
534 "2023-01-01",
535 "--cutoff-date",
536 "2023-04-01",
537 "--transfer-list",
538 "https://raw.githubusercontent.com/test/cm_transfer_list.yaml",
539 "--dry-run",
540 "--out-dir",
541 self.out_dir,
542 ],
543 )
545 self.assertEqual(result.exit_code, 0, msg=result.output)
546 sanitized_repo = self.repo.strip("/").replace("/", "_")
547 expected_dir = os.path.join(self.out_dir, sanitized_repo, "auto_register")
548 self.assertTrue(os.path.exists(expected_dir))
549 files = os.listdir(expected_dir)
550 self.assertGreater(len(files), 0)
552 with open(os.path.join(expected_dir, files[0])) as f:
553 content = f.read()
555 self.assertIn("22222222-2222-2222-2222-222222222222", content)
556 self.assertNotIn("11111111-1111-1111-1111-111111111111", content)
557 self.assertNotIn("33333333-3333-3333-3333-333333333333", content)
559 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor)
560 @patch("lsst.rucio.register.script._get_rucio_interface")
561 @patch("lsst.rucio.register.script.Butler")
562 def test_auto_register_with_dataset_name_prefix(self, mock_butler_class, mock_get_ri):
563 mock_butler = MagicMock()
564 mock_butler_class.return_value = mock_butler
565 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"]
567 mock_info = MagicMock()
568 mock_info.dataset_types = ["typeA"]
569 mock_butler.collections.get_info.return_value = mock_info
571 mock_ri = MagicMock()
572 mock_inner_butler = MagicMock()
573 mock_get_ri.return_value = (mock_ri, mock_inner_butler)
575 refA = MagicMock()
576 refA.id = "11111111-1111-1111-1111-111111111111"
577 refA.dataset_type.name = "typeA"
579 mock_butler.query_all_datasets.return_value = [refA]
581 runner = CliRunner()
582 result = runner.invoke(
583 script.main,
584 [
585 "auto-register",
586 "--repo",
587 self.repo,
588 "--root-chain",
589 "HSC/runs/*",
590 "--start-date",
591 "2023-01-01",
592 "--cutoff-date",
593 "2023-04-01",
594 "--dataset-name-prefix",
595 "Dataset/HSC/runs/RC2/w_2026_23/DM-55175/LANCS/outputs",
596 "--dry-run",
597 "--out-dir",
598 self.out_dir,
599 ],
600 )
602 self.assertEqual(result.exit_code, 0, msg=result.output)
603 sanitized_repo = self.repo.strip("/").replace("/", "_")
604 expected_dir = os.path.join(self.out_dir, sanitized_repo, "auto_register")
605 self.assertTrue(os.path.exists(expected_dir))
606 files = os.listdir(expected_dir)
607 self.assertTrue(
608 any("Dataset-HSC-runs-RC2-w_2026_23-DM-55175-LANCS-outputs-typeA-USDF" in f for f in files)
609 )
611 @patch("lsst.rucio.register.script.ProcessPoolExecutor", DummyExecutor)
612 @patch("lsst.rucio.register.script._process_auto_register_batch")
613 @patch("lsst.rucio.register.script.Butler")
614 def test_auto_register_with_clear_is_new(self, mock_butler_class, mock_process_batch):
615 mock_butler = MagicMock()
616 mock_butler_class.return_value = mock_butler
617 mock_process_batch.return_value = {"registered": 1, "failed": 0}
619 mock_butler.collections.query.return_value = ["HSC/runs/RC2/w_2023_32"]
620 mock_info = MagicMock()
621 mock_info.dataset_types = ["typeA"]
622 mock_butler.collections.get_info.return_value = mock_info
624 refA = MagicMock()
625 refA.id = "11111111-1111-1111-1111-111111111111"
626 refA.dataset_type.name = "typeA"
627 mock_butler.query_all_datasets.return_value = [refA]
629 runner = CliRunner()
630 result = runner.invoke(
631 script.main,
632 [
633 "auto-register",
634 "--repo",
635 self.repo,
636 "--root-chain",
637 "HSC/runs/*",
638 "--start-date",
639 "2023-01-01",
640 "--cutoff-date",
641 "2023-04-01",
642 "--clear-is-new",
643 ],
644 )
646 self.assertEqual(result.exit_code, 0, msg=result.output)
647 self.assertTrue(mock_process_batch.called)
648 _, kwargs = mock_process_batch.call_args
649 self.assertTrue(kwargs.get("clear_is_new"))
651 @patch("lsst.rucio.register.script._get_rucio_interface")
652 def test_dataset_list_with_json_uuidlist(self, mock_get_ri):
653 mock_ri = MagicMock()
654 mock_butler = MagicMock()
655 mock_get_ri.return_value = (mock_ri, mock_butler)
657 json_file = os.path.join(self.tmp_dir.name, "auto-register-failures.json")
658 with open(json_file, "w") as f:
659 json.dump(["019ee079-b7e4-700d-a0bf-03412ae4fe0a"], f)
661 mock_ref = MagicMock()
662 mock_ref.id = "019ee079-b7e4-700d-a0bf-03412ae4fe0a"
663 mock_butler.get_many_datasets.return_value = [mock_ref]
665 runner = CliRunner()
666 result = runner.invoke(
667 script.main,
668 [
669 "dataset-list",
670 "--repo",
671 self.repo,
672 "--rucio-dataset",
673 "Dataset/Test/Failed",
674 "--uuidlist",
675 json_file,
676 ],
677 )
679 self.assertEqual(result.exit_code, 0, msg=result.output)
680 mock_butler.get_many_datasets.assert_called_once_with(["019ee079-b7e4-700d-a0bf-03412ae4fe0a"])
683class MemoryTester(lsst.utils.tests.MemoryTestCase):
684 pass
687def setup_module(module):
688 lsst.utils.tests.init()