Coverage for tests/test_interface.py: 100%

184 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-11-07 01:39 -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/>. 

21 

22 

23import os 

24import shutil 

25import tempfile 

26from unittest.mock import MagicMock, patch 

27 

28import urllib3 

29from rucio.client.didclient import DIDClient 

30from rucio.client.replicaclient import ReplicaClient 

31from rucio.common.exception import ( 

32 DataIdentifierNotFound, 

33 FileAlreadyExists, 

34 RucioException, 

35) 

36 

37import lsst.utils.tests 

38from lsst.daf.butler import Butler, DatasetRef, DimensionUniverse 

39from lsst.resources import ResourceInfo, ResourcePath 

40from lsst.resources.file import FileResourcePath 

41from lsst.rucio.register.data_type import DataType 

42from lsst.rucio.register.rucio_interface import RucioInterface 

43 

44 

45class InterfaceTestCase(lsst.utils.tests.TestCase): 

46 maxDiff = None 

47 FAST_BACKOFF = {"factor": 0.001, "max_value": 0.01, "max_tries": 2} 

48 

49 def setUp(self): 

50 self.butler_repo = tempfile.mkdtemp(dir="/tmp") 

51 test_dir = os.path.abspath(os.path.dirname(__file__)) 

52 

53 self.dataset_ref_file = os.path.join(test_dir, "data", "dataset_ref.json") 

54 

55 Butler.makeRepo(self.butler_repo) 

56 

57 data_name = "visitSummary_HSC_y_HSC-Y_318_HSC_runs_RC2_w_2023_32_DM-40356_20230814T170253Z.fits" 

58 json_name = "visitSummary_HSC_y_HSC-Y_318_HSC_runs_RC2_w_2023_32_DM-40356_20230814T170253Z.json" 

59 

60 self.data_file = os.path.join(test_dir, "data", data_name) 

61 self.json_file = os.path.join(test_dir, "data", json_name) 

62 

63 self.butler = Butler(self.butler_repo, writeable=True) 

64 self.butler.getURI = MagicMock(return_value=ResourcePath(f"file://{self.data_file}")) 

65 

66 self.rse_root = tempfile.mkdtemp(dir="/tmp") 

67 

68 # patch __init__ methods 

69 self.rc_init = patch.object(ReplicaClient, "__init__", return_value=None) 

70 self.dc_init = patch.object(DIDClient, "__init__", return_value=None) 

71 self.rc_add_replicas = patch.object(ReplicaClient, "add_replicas", return_value=None) 

72 self.dc_attach_dids = patch.object(DIDClient, "attach_dids", return_value=None) 

73 self.dc_attach_dids_to_dids = patch.object(DIDClient, "attach_dids_to_dids", return_value=None) 

74 self.dc_add_dataset = patch.object(DIDClient, "add_dataset", return_value=None) 

75 self.rand = patch("random.randint", return_value=1) 

76 

77 self.mock_rc_init = self.rc_init.start() 

78 self.mock_dc_init = self.dc_init.start() 

79 self.mock_rc_add_replicas = self.rc_add_replicas.start() 

80 self.mock_dc_attach_dids = self.dc_attach_dids.start() 

81 self.mock_dc_attach_dids_to_dids = self.dc_attach_dids_to_dids.start() 

82 self.mock_dc_add_dataset = self.dc_add_dataset.start() 

83 self.mock_rand = self.rand.start() 

84 

85 rucio_rse = "DRR1" 

86 scope = "test" 

87 dtn_url = "root://xrd1:1094//rucio" 

88 self.ri = RucioInterface(self.butler, rucio_rse, scope, self.rse_root, dtn_url, DataType.DATA_PRODUCT) 

89 self.rpath = ResourcePath(f"{self.data_file}") 

90 self.fpath = FileResourcePath(f"{self.data_file}") 

91 

92 def testResourcePathCase(self): 

93 res = self.ri.compute_hashes(self.rpath) 

94 self.assertEqual(res, (1365120, "480be4de")) 

95 

96 def testFileResourcePathCase(self): 

97 res = self.ri.compute_hashes(self.fpath) 

98 self.assertEqual(res, (1365120, "480be4de")) 

99 

100 def testChecksumsCase(self): 

101 fake_info = MagicMock(spec=ResourceInfo) 

102 fake_info.checksums = {"adler32": "abcd1234"} 

103 fake_info.size = 1234 

104 fake_info.is_file = True 

105 fake_info.last_modified = None 

106 fake_info.uri = f"file://{self.data_file}" 

107 

108 patcher = patch.object(FileResourcePath, "get_info", return_value=fake_info) 

109 patcher.start() 

110 self.addCleanup(patcher.stop) 

111 

112 with open(self.dataset_ref_file) as f: 

113 json_ref = f.readline() 

114 

115 ref = DatasetRef.from_json(json_ref, DimensionUniverse()) 

116 

117 self.butler.registry.registerDatasetType(ref.datasetType) 

118 cnt = self.ri.register_as_replicas("mydataset", [ref]) 

119 self.assertEqual(cnt, 1) 

120 

121 rb = self.ri._make_dataset_ref_bundle("mydataset", ref) 

122 self.assertEqual(rb.dataset_id, "mydataset") 

123 

124 did = rb.did.model_dump() 

125 self.assertEqual(did["adler32"], "abcd1234") 

126 

127 def testHashCacheReusesComputedChecksum(self): 

128 fake_path = MagicMock() 

129 fake_path.scheme = "custom" 

130 fake_path.__str__.return_value = "custom://repo/file.fits" 

131 fake_path.get_info.return_value = MagicMock(size=1234, checksums={"adler32": "12345678"}) 

132 

133 res1 = self.ri.compute_hashes(fake_path) 

134 self.assertEqual(res1, (1234, "12345678")) 

135 self.assertEqual(self.ri._checksum_sources["standard WebDAV metadata"], 1) 

136 

137 res2 = self.ri.compute_hashes(fake_path) 

138 self.assertEqual(res2, (1234, "12345678")) 

139 self.assertEqual(self.ri._checksum_sources["internal hash cache"], 1) 

140 

141 def testInterfaceTestCase(self): 

142 dtn_url = "root://xrd1:1094//rucio" 

143 

144 json_ref = None 

145 with open(self.dataset_ref_file) as f: 

146 json_ref = f.readline() 

147 

148 ref = DatasetRef.from_json(json_ref, DimensionUniverse()) 

149 

150 self.butler.registry.registerDatasetType(ref.datasetType) 

151 cnt = self.ri.register_as_replicas("mydataset", [ref]) 

152 self.assertEqual(cnt, 1) 

153 

154 rb = self.ri._make_dataset_ref_bundle("mydataset", ref) 

155 self.assertEqual(rb.dataset_id, "mydataset") 

156 

157 did = rb.did.model_dump() 

158 self.assertEqual(did["pfn"], f"{dtn_url}{self.data_file}") 

159 self.assertEqual(did["bytes"], 1365120) 

160 self.assertEqual(did["adler32"], "480be4de") 

161 self.assertEqual(did["name"], self.data_file) 

162 self.assertEqual(did["scope"], "test") 

163 

164 meta = did["meta"] 

165 self.assertEqual(meta["rubin_butler"], DataType.DATA_PRODUCT) 

166 

167 def common(self): 

168 json_ref = None 

169 with open(self.dataset_ref_file) as f: 

170 json_ref = f.readline() 

171 

172 ref = DatasetRef.from_json(json_ref, DimensionUniverse()) 

173 

174 self.butler.registry.registerDatasetType(ref.datasetType) 

175 self.ri.register_as_replicas("mydataset", [ref]) 

176 

177 @patch.dict("lsst.rucio.register.rucio_interface._BACKOFF", FAST_BACKOFF) 

178 @patch.object(ReplicaClient, "add_replicas", side_effect=RucioException("failed")) 

179 def testException1TestCase(self, MC1): 

180 self.ri.register_to_dataset = MagicMock(name="register_to_dataset") 

181 with self.assertRaises(Exception): 

182 self.common() 

183 

184 @patch.dict("lsst.rucio.register.rucio_interface._BACKOFF", FAST_BACKOFF) 

185 @patch.object(DIDClient, "add_files_to_dataset", side_effect=FileAlreadyExists("failed")) 

186 def testException2TestCase(self, MC1): 

187 self.common() 

188 

189 def testClearIsNewMetadata(self): 

190 self.ri.clear_is_new = True 

191 self.ri.did_client = MagicMock() 

192 self.ri.did_client.set_metadata = MagicMock(name="set_metadata") 

193 

194 fake_bundle = MagicMock() 

195 fake_bundle.dataset_id = "dataset_001" 

196 fake_bundle.get_did.return_value = {"pfn": "test_pfn"} 

197 

198 with patch.object(self.ri, "_add_files_to_dataset"): 

199 self.ri.register_to_dataset([fake_bundle]) 

200 

201 self.ri.did_client.set_metadata.assert_called_once_with( 

202 scope="test", name="dataset_001", key="is_new", value=None 

203 ) 

204 

205 def testRetriesConfigured(self): 

206 self.ri.set_backoff(factor=2.0, max_value=60, max_tries=12) 

207 self.assertEqual(self.ri.retries, 12) 

208 

209 def testTimeoutConfigured(self): 

210 self.assertEqual(self.ri.timeout, RucioInterface.DEFAULT_TIMEOUT) 

211 self.assertEqual(RucioInterface.DEFAULT_TIMEOUT, 300.0) 

212 self.mock_rc_init.assert_called_with(timeout=300.0) 

213 self.mock_dc_init.assert_called_with(timeout=300.0) 

214 

215 custom_ri = RucioInterface( 

216 self.butler, 

217 "DRR1", 

218 "test", 

219 self.rse_root, 

220 "root://xrd1:1094//rucio", 

221 DataType.DATA_PRODUCT, 

222 timeout=120.0, 

223 ) 

224 self.assertEqual(custom_ri.timeout, 120.0) 

225 self.mock_rc_init.assert_called_with(timeout=120.0) 

226 self.mock_dc_init.assert_called_with(timeout=120.0) 

227 

228 @patch.dict("lsst.rucio.register.rucio_interface._BACKOFF", FAST_BACKOFF) 

229 @patch.object( 

230 ReplicaClient, 

231 "add_replicas", 

232 side_effect=urllib3.exceptions.ReadTimeoutError(pool=None, url="http://127.0.0.1/", message="failed"), 

233 ) 

234 def testException3TestCase(self, MC3): 

235 self.ri.register_to_dataset = MagicMock(name="register_to_dataset") 

236 with self.assertRaises(Exception): 

237 self.common() 

238 

239 @patch.dict("lsst.rucio.register.rucio_interface._BACKOFF", FAST_BACKOFF) 

240 @patch.object( 

241 DIDClient, 

242 "add_files_to_datasets", 

243 side_effect=urllib3.exceptions.ReadTimeoutError(pool=None, url="http://127.0.0.1/", message="failed"), 

244 ) 

245 def testException4TestCase(self, MC4): 

246 with self.assertRaises(Exception): 

247 self.common() 

248 

249 @patch.dict("lsst.rucio.register.rucio_interface._BACKOFF", FAST_BACKOFF) 

250 @patch.object(DIDClient, "add_files_to_datasets", side_effect=DataIdentifierNotFound("failed")) 

251 def testException5TestCase(self, MC5): 

252 with self.assertRaises(Exception): 

253 self.common() 

254 

255 @patch.dict("lsst.rucio.register.rucio_interface._BACKOFF", FAST_BACKOFF) 

256 @patch.object(DIDClient, "add_files_to_datasets", side_effect=DataIdentifierNotFound("failed")) 

257 @patch.object( 

258 DIDClient, 

259 "add_dataset", 

260 side_effect=urllib3.exceptions.ReadTimeoutError(pool=None, url="http://127.0.0.1/", message="failed"), 

261 ) 

262 def testException6TestCase(self, MC6a, MC6b): 

263 with self.assertRaises(Exception): 

264 self.common() 

265 

266 def testMultiThreadedRegisterAsReplicas(self): 

267 json_ref = None 

268 with open(self.dataset_ref_file) as f: 

269 json_ref = f.readline() 

270 

271 ref = DatasetRef.from_json(json_ref, DimensionUniverse()) 

272 self.butler.registry.registerDatasetType(ref.datasetType) 

273 

274 refs = [ref] * 25 

275 cnt = self.ri.register_as_replicas("multithread_dataset", refs) 

276 self.assertEqual(cnt, 25) 

277 self.mock_rc_add_replicas.assert_called() 

278 

279 def tearDown(self): 

280 patch.stopall() 

281 shutil.rmtree(self.butler_repo, ignore_errors=True) 

282 shutil.rmtree(self.rse_root, ignore_errors=True) 

283 

284 

285class MemoryTester(lsst.utils.tests.MemoryTestCase): 

286 pass 

287 

288 

289def setup_module(module): 

290 lsst.utils.tests.init()