Coverage for tests/test_s3.py: 29%

1# This file is part of lsst-resources. 

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# Use of this source code is governed by a 3-clause BSD-style 

10# license that can be found in the LICENSE file. 

11 

12import os 

13import time 

14import unittest 

15from inspect import signature 

16from unittest import mock 

17from urllib.parse import parse_qs, urlparse 

18 

19from lsst.resources import ResourcePath 

20from lsst.resources.s3 import S3ResourcePath 

21from lsst.resources.s3utils import clean_test_environment_for_s3 

22from lsst.resources.tests import GenericReadWriteTestCase, GenericTestCase 

23 

24try: 

25 import boto3 

26 import botocore 

27 

28 try: 

29 from moto import mock_aws # v5 

30 except ImportError: 

31 from moto import mock_s3 as mock_aws 

32except ImportError: 

33 boto3 = None 

34 

35 def mock_aws(cls): 

36 """No-op decorator in case moto mock_aws can not be imported.""" 

37 return cls 

38 

39 

40try: 

41 import fsspec 

42except ImportError: 

43 fsspec = None 

44 

45 

46class GenericS3TestCase(GenericTestCase, unittest.TestCase): 

47 """Generic tests of S3 URIs.""" 

48 

49 scheme = "s3" 

50 netloc = "my_bucket" 

51 

52 

53class S3ReadWriteTestCaseBase(GenericReadWriteTestCase): 

54 """Tests of reading and writing S3 URIs.""" 

55 

56 scheme = "s3" 

57 s3_endpoint_url: str | None = None 

58 # S3 batches in 1000 files so need more than that. 

59 n_mremove_files: int = 1015 

60 

61 def setUp(self): 

62 self.enterContext(clean_test_environment_for_s3()) 

63 

64 # Enable S3 mocking of tests. 

65 self.enterContext(mock_aws()) 

66 

67 # MOTO needs to know that we expect Bucket bucketname to exist 

68 s3 = boto3.resource("s3", endpoint_url=self.s3_endpoint_url) 

69 s3.create_bucket(Bucket=self.bucket) 

70 

71 super().setUp() 

72 

73 def tearDown(self): 

74 s3 = boto3.resource("s3") 

75 bucket = s3.Bucket(self.bucket) 

76 try: 

77 bucket.objects.all().delete() 

78 except botocore.exceptions.ClientError as e: 

79 if e.response["Error"]["Code"] == "404": 

80 # the key was not reachable - pass 

81 pass 

82 else: 

83 raise 

84 

85 bucket = s3.Bucket(self.bucket) 

86 bucket.delete() 

87 

88 S3ResourcePath.use_threads = None 

89 

90 super().tearDown() 

91 

92 def test_bucket_fail(self): 

93 # Deliberately create URI with unknown bucket. 

94 uri = ResourcePath("s3://badbucket/something/") 

95 

96 with self.assertRaises(ValueError): 

97 uri.mkdir() 

98 

99 with self.assertRaises(FileNotFoundError): 

100 uri.remove() 

101 

102 def test_transfer_progress(self): 

103 """Test progress bar reporting for upload and download.""" 

104 remote = self.root_uri.join("test.dat") 

105 remote.write(b"42") 

106 with ResourcePath.temporary_uri(suffix=".dat") as tmp: 

107 # Download from S3. 

108 with self.assertLogs("lsst.resources", level="DEBUG") as cm: 

109 tmp.transfer_from(remote, transfer="auto") 

110 self.assertRegex("".join(cm.output), r"test\.dat.*100\%") 

111 

112 # Upload to S3. 

113 with self.assertLogs("lsst.resources", level="DEBUG") as cm: 

114 remote.transfer_from(tmp, transfer="auto", overwrite=True) 

115 self.assertRegex("".join(cm.output), rf"{tmp.basename()}.*100\%") 

116 

117 def test_handle(self): 

118 remote = self.root_uri.join("test_handle.dat") 

119 with remote.open("wb") as handle: 

120 self.assertTrue(handle.writable()) 

121 # write 6 megabytes to make sure partial write work 

122 handle.write(6 * 1024 * 1024 * b"a") 

123 self.assertEqual(handle.tell(), 6 * 1024 * 1024) 

124 handle.flush() 

125 self.assertGreaterEqual(len(handle._multiPartUpload), 1) 

126 

127 # verify file can't be seeked back 

128 with self.assertRaises(OSError): 

129 handle.seek(0) 

130 

131 # write more bytes 

132 handle.write(1024 * b"c") 

133 

134 # seek back and overwrite 

135 handle.seek(6 * 1024 * 1024) 

136 handle.write(1024 * b"b") 

137 

138 with remote.open("rb") as handle: 

139 self.assertTrue(handle.readable()) 

140 # read the first 6 megabytes 

141 result = handle.read(6 * 1024 * 1024) 

142 self.assertEqual(result, 6 * 1024 * 1024 * b"a") 

143 self.assertEqual(handle.tell(), 6 * 1024 * 1024) 

144 # verify additional read gets the next part 

145 result = handle.read(1024) 

146 self.assertEqual(result, 1024 * b"b") 

147 # see back to the beginning to verify seeking 

148 handle.seek(0) 

149 result = handle.read(1024) 

150 self.assertEqual(result, 1024 * b"a") 

151 

152 remote = self.root_uri.join("missing_file.dat") 

153 with remote.open("rb") as handle: 

154 with self.assertRaises(FileNotFoundError): 

155 handle.read() 

156 

157 def test_url_signing(self): 

158 self._test_url_signing_case("url-signing-test.txt", b"test123") 

159 # A zero byte presigned S3 HTTP URL is a weird edge case, because we 

160 # emulate HEAD requests using a 1-byte GET. 

161 self._test_url_signing_case("url-signing-test-zero-bytes.txt", b"") 

162 # Should be the same as a normal case, but check it for paranoia since 

163 # it's on the boundary of the read size. 

164 self._test_url_signing_case("url-signing-test-one-byte.txt", b"t") 

165 

166 def _test_url_signing_case(self, filename: str, test_data: bytes): 

167 s3_path = self.root_uri.join(filename) 

168 

169 put_url = s3_path.generate_presigned_put_url(expiration_time_seconds=1800) 

170 self._check_presigned_url(put_url, 1800) 

171 get_url = s3_path.generate_presigned_get_url(expiration_time_seconds=3600) 

172 self._check_presigned_url(get_url, 3600) 

173 

174 # Check that fragments are retained. 

175 s3_path = s3_path.replace(fragment="zip-path=X") 

176 put_url = s3_path.generate_presigned_put_url(expiration_time_seconds=1800) 

177 self.assertEqual(ResourcePath(put_url).fragment, "zip-path=X") 

178 self._check_presigned_url(put_url, 1800) 

179 get_url = s3_path.generate_presigned_get_url(expiration_time_seconds=3600) 

180 self.assertEqual(ResourcePath(get_url).fragment, "zip-path=X") 

181 self._check_presigned_url(get_url, 3600) 

182 

183 # Moto monkeypatches the 'requests' library to mock access to presigned 

184 # URLs, so we are able to use HttpResourcePath to access the URLs in 

185 # this test. 

186 ResourcePath(put_url).write(test_data) 

187 get_path = ResourcePath(get_url) 

188 retrieved = get_path.read() 

189 self.assertEqual(retrieved, test_data) 

190 self.assertTrue(get_path.exists()) 

191 self.assertEqual(get_path.size(), len(test_data)) 

192 

193 # Try again with open(). 

194 with get_path.open("rb") as fd: 

195 self.assertEqual(fd.read(), test_data) 

196 

197 def test_nonexistent_presigned_url(self): 

198 s3_path = self.root_uri.join("this-is-a-missing-file.txt") 

199 get_url = s3_path.generate_presigned_get_url(expiration_time_seconds=3600) 

200 get_path = ResourcePath(get_url) 

201 # Check the HttpResourcePath implementation for presigned S3 urls. 

202 # Nothing has been uploaded to this URL, so it shouldn't exist. 

203 self.assertFalse(get_path.exists()) 

204 with self.assertRaises(FileNotFoundError): 

205 get_path.size() 

206 

207 def _check_presigned_url(self, url: str, expiration_time_seconds: int): 

208 parsed = urlparse(url) 

209 self.assertEqual(parsed.scheme, "https") 

210 

211 actual_expiration_timestamp = int(parse_qs(parsed.query)["Expires"][0]) 

212 current_time = int(time.time()) 

213 expected_expiration_timestamp = current_time + expiration_time_seconds 

214 # Allow some flex in the expiration time in case this test process goes 

215 # out to lunch for a while on a busy CI machine 

216 self.assertLessEqual(abs(expected_expiration_timestamp - actual_expiration_timestamp), 120) 

217 

218 def test_threading_true(self): 

219 with mock.patch.dict(os.environ, {"LSST_S3_USE_THREADS": "True"}): 

220 S3ResourcePath.use_threads = None 

221 test_resource_path = self.root_uri.join("test_file.dat") 

222 self.assertTrue(test_resource_path._transfer_config.use_threads) 

223 

224 def test_implicit_default_threading(self): 

225 S3ResourcePath.use_threads = None 

226 boto_default = signature(boto3.s3.transfer.TransferConfig).parameters["use_threads"].default 

227 test_resource_path = self.root_uri.join("test_file.dat") 

228 self.assertEqual(test_resource_path._transfer_config.use_threads, boto_default) 

229 

230 def test_explicit_default_threading(self): 

231 with mock.patch.dict(os.environ, {"LSST_S3_USE_THREADS": "None"}): 

232 S3ResourcePath.use_threads = None 

233 boto_default = signature(boto3.s3.transfer.TransferConfig).parameters["use_threads"].default 

234 test_resource_path = self.root_uri.join("test_file.dat") 

235 self.assertEqual(test_resource_path._transfer_config.use_threads, boto_default) 

236 

237 def test_threading_false(self): 

238 with mock.patch.dict(os.environ, {"LSST_S3_USE_THREADS": "False"}): 

239 S3ResourcePath.use_threads = None 

240 test_resource_path = self.root_uri.join("test_file.dat") 

241 self.assertFalse(test_resource_path._transfer_config.use_threads) 

242 

243 self.test_local() 

244 

245 @unittest.skipIf(fsspec is None, "fsspec is not available") 

246 def test_fsspec_constructor(self) -> None: 

247 """Test that we can obtain an s3fs object.""" 

248 uri = self.root_uri.join("test_file.dat") 

249 fs, path = uri.to_fsspec() 

250 self.assertEqual(path, f"{uri._bucket}/{uri.relativeToPathRoot}") 

251 self.assertTrue(hasattr(fs, "open")) 

252 

253 def test_fsspec(self) -> None: 

254 raise unittest.SkipTest("fsspec s3fs incompatible with moto") 

255 

256 @unittest.mock.patch("lsst.resources._resourcePath._POOL_EXECUTOR_CLASS", None) 

257 @unittest.mock.patch.dict(os.environ, {"LSST_RESOURCES_EXECUTOR": "threads"}) 

258 def test_mexists(self) -> None: 

259 """Test mexists with override executor pool. 

260 

261 moto does not work with process pool. 

262 """ 

263 super().test_mexists() 

264 

265 @unittest.mock.patch("lsst.resources._resourcePath._POOL_EXECUTOR_CLASS", None) 

266 @unittest.mock.patch.dict(os.environ, {"LSST_RESOURCES_EXECUTOR": "threads"}) 

267 def test_mtransfer(self) -> None: 

268 """Test mtransfer with override executor pool. 

269 

270 moto does not work with process pool. 

271 """ 

272 super().test_mtransfer() 

273 

274 @unittest.mock.patch("lsst.resources._resourcePath._POOL_EXECUTOR_CLASS", None) 

275 @unittest.mock.patch.dict(os.environ, {"LSST_RESOURCES_EXECUTOR": "threads"}) 

276 def test_local_mtransfer(self) -> None: 

277 """Test local mtransfer with override executor pool. 

278 

279 moto does not work with process pool. 

280 """ 

281 super().test_local_mtransfer() 

282 

283 

284@unittest.skipIf(not boto3, "Warning: boto3 AWS SDK not found!") 

285class S3ReadWriteTestCase(S3ReadWriteTestCaseBase, unittest.TestCase): 

286 """Test S3 with no explicit profile/endpoint specified. 

287 (``s3://bucketname/...``). 

288 """ 

289 

290 bucket = "my_2nd_bucket" 

291 netloc = bucket 

292 

293 

294@unittest.skipIf(not boto3, "Warning: boto3 AWS SDK not found!") 

295class S3WithProfileReadWriteTestCase(S3ReadWriteTestCaseBase, unittest.TestCase): 

296 """Test S3 URLs with explicit profile specified. 

297 (``s3://profile@bucketname/...``). 

298 """ 

299 

300 bucket = "3rd_bucket" 

301 netloc = f"myprofile@{bucket}" 

302 s3_endpoint_url = "https://endpoint1.test.example" 

303 

304 def setUp(self): 

305 # Configure custom S3 endpoints that we can target from tests using 

306 # non-default profile. 

307 self.enterContext( 

308 mock.patch.dict( 

309 os.environ, 

310 { 

311 "MOTO_S3_CUSTOM_ENDPOINTS": self.s3_endpoint_url, 

312 "LSST_RESOURCES_S3_PROFILE_myprofile": "https://access_key:security_key@endpoint1.test.example", 

313 }, 

314 ) 

315 ) 

316 

317 super().setUp() 

318 

319 def test_missing_profile(self): 

320 with self.assertRaises(botocore.exceptions.ProfileNotFound): 

321 ResourcePath("s3://otherprofile@bucket").read() 

322 

323 def test_s3_endpoint_url(self): 

324 with mock.patch.dict( 

325 os.environ, 

326 {"S3_ENDPOINT_URL": self.s3_endpoint_url}, 

327 ): 

328 path = ResourcePath(f"s3://{self.bucket}/test-s3-endpoint-url.txt") 

329 data = b"123" 

330 path.write(data) 

331 self.assertEqual(path.read(), data) 

332 self.assertIn( 

333 "https://endpoint1.test.example", 

334 path.generate_presigned_get_url(expiration_time_seconds=3600), 

335 ) 

336 

337 def test_uri_syntax(self): 

338 path1 = ResourcePath("s3://profile@bucket/path") 

339 self.assertEqual(path1._bucket, "bucket") 

340 self.assertEqual(path1._profile, "profile") 

341 path2 = ResourcePath("s3://bucket2/path") 

342 self.assertEqual(path2._bucket, "bucket2") 

343 self.assertIsNone(path2._profile) 

344 

345 def test_ceph_uri_syntax(self): 

346 # The Ceph S3 'multi-tenant' syntax for buckets can include colons. 

347 path1 = ResourcePath("s3://profile@ceph:bucket/path") 

348 self.assertEqual(path1._bucket, "ceph:bucket") 

349 self.assertEqual(path1._profile, "profile") 

350 path2 = ResourcePath("s3://ceph:bucket2/path") 

351 self.assertEqual(path2._bucket, "ceph:bucket2") 

352 self.assertIsNone(path2._profile) 

353 

354 def test_transfer_from_different_endpoints(self): 

355 # Create a bucket using a different endpoint (the default endpoint.) 

356 boto3.resource("s3").create_bucket(Bucket="source-bucket") 

357 source_path = ResourcePath("s3://source-bucket/file.txt") 

358 source_path.write(b"123") 

359 target_path = ResourcePath(f"s3://{self.netloc}/target.txt") 

360 # Transfer from default endpoint to custom endpoint with custom 

361 # profile. 

362 target_path.transfer_from(source_path) 

363 self.assertEqual(target_path.read(), b"123") 

364 

365 

366if __name__ == "__main__": 

367 unittest.main()