Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

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

21 

22import os 

23import unittest 

24import tempfile 

25import shutil 

26import string 

27import random 

28 

29try: 

30 import boto3 

31 import botocore 

32 from moto import mock_s3 

33except ImportError: 

34 boto3 = None 

35 

36 def mock_s3(cls): 

37 """A no-op decorator in case moto mock_s3 can not be imported. 

38 """ 

39 return cls 

40 

41import lsst.utils.tests 

42 

43from lsst.daf.butler import Butler, Config 

44from lsst.daf.butler import StorageClassFactory 

45from lsst.daf.butler import DatasetType 

46from lsst.daf.butler.core.location import ButlerURI 

47from lsst.daf.butler.core.s3utils import setAwsEnvCredentials, unsetAwsEnvCredentials 

48from lsst.daf.butler.tests import FitsCatalogDatasetsHelper, DatasetTestHelper 

49 

50try: 

51 import lsst.afw.image 

52 from lsst.afw.image import LOCAL 

53 from lsst.geom import Box2I, Point2I, Extent2I 

54except ImportError: 

55 lsst.afw = None 

56 

57TESTDIR = os.path.dirname(__file__) 

58 

59 

60class ButlerFitsTests(FitsCatalogDatasetsHelper, DatasetTestHelper): 

61 useTempRoot = True 

62 

63 @staticmethod 

64 def registerDatasetTypes(datasetTypeName, dimensions, storageClass, registry): 

65 """Bulk register DatasetTypes 

66 """ 

67 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

68 registry.registerDatasetType(datasetType) 

69 

70 for compName, compStorageClass in storageClass.components.items(): 

71 compType = DatasetType(datasetType.componentTypeName(compName), dimensions, compStorageClass) 

72 registry.registerDatasetType(compType) 

73 

74 @classmethod 

75 def setUpClass(cls): 

76 if lsst.afw is None: 

77 raise unittest.SkipTest("afw not available.") 

78 cls.storageClassFactory = StorageClassFactory() 

79 cls.storageClassFactory.addFromConfig(cls.configFile) 

80 

81 def setUp(self): 

82 """Create a new butler root for each test.""" 

83 if self.useTempRoot: 

84 self.root = tempfile.mkdtemp(dir=TESTDIR) 

85 Butler.makeRepo(self.root, config=Config(self.configFile)) 

86 self.tmpConfigFile = os.path.join(self.root, "butler.yaml") 

87 else: 

88 self.root = None 

89 self.tmpConfigFile = self.configFile 

90 

91 def tearDown(self): 

92 if self.root is not None and os.path.exists(self.root): 

93 shutil.rmtree(self.root, ignore_errors=True) 

94 

95 def testExposureCompositePutGetConcrete(self): 

96 storageClass = self.storageClassFactory.getStorageClass("ExposureF") 

97 self.runExposureCompositePutGetTest(storageClass, "calexp") 

98 

99 def testExposureCompositePutGetVirtual(self): 

100 storageClass = self.storageClassFactory.getStorageClass("ExposureCompositeF") 

101 self.runExposureCompositePutGetTest(storageClass, "unknown") 

102 

103 def runExposureCompositePutGetTest(self, storageClass, datasetTypeName): 

104 example = os.path.join(TESTDIR, "data", "basic", "small.fits") 

105 exposure = lsst.afw.image.ExposureF(example) 

106 butler = Butler(self.tmpConfigFile, run="ingest") 

107 dimensions = butler.registry.dimensions.extract(["instrument", "visit"]) 

108 self.registerDatasetTypes(datasetTypeName, dimensions, storageClass, butler.registry) 

109 dataId = {"visit": 42, "instrument": "DummyCam", "physical_filter": "d-r"} 

110 # Add needed Dimensions 

111 butler.registry.insertDimensionData("instrument", {"instrument": "DummyCam"}) 

112 butler.registry.insertDimensionData("physical_filter", {"instrument": "DummyCam", "name": "d-r", 

113 "abstract_filter": "R"}) 

114 butler.registry.insertDimensionData("visit", {"instrument": "DummyCam", "id": 42, 

115 "name": "fortytwo", "physical_filter": "d-r"}) 

116 butler.put(exposure, datasetTypeName, dataId) 

117 # Get the full thing 

118 butler.get(datasetTypeName, dataId) 

119 # TODO enable check for equality (fix for Exposure type) 

120 # self.assertEqual(full, exposure) 

121 # Get a component 

122 compsRead = {} 

123 for compName in ("wcs", "image", "mask", "coaddInputs", "psf"): 

124 compTypeName = DatasetType.nameWithComponent(datasetTypeName, compName) 

125 component = butler.get(compTypeName, dataId) 

126 # TODO enable check for component instance types 

127 # compRef = butler.registry.find(butler.run.name, 

128 # f"calexp.{compName}", dataId) 

129 # self.assertIsInstance(component, 

130 # compRef.datasetType.storageClass.pytype) 

131 compsRead[compName] = component 

132 # Simple check of WCS 

133 bbox = Box2I(Point2I(0, 0), Extent2I(9, 9)) 

134 self.assertWcsAlmostEqualOverBBox(compsRead["wcs"], exposure.getWcs(), bbox) 

135 

136 # With parameters 

137 inBBox = Box2I(minimum=Point2I(0, 0), maximum=Point2I(3, 3)) 

138 parameters = dict(bbox=inBBox, origin=LOCAL) 

139 subset = butler.get(datasetTypeName, dataId, parameters=parameters) 

140 outBBox = subset.getBBox() 

141 self.assertEqual(inBBox, outBBox) 

142 

143 

144class PosixDatastoreButlerTestCase(ButlerFitsTests, lsst.utils.tests.TestCase): 

145 """PosixDatastore specialization of a butler""" 

146 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

147 

148 

149class InMemoryDatastoreButlerTestCase(ButlerFitsTests, lsst.utils.tests.TestCase): 

150 """InMemoryDatastore specialization of a butler""" 

151 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml") 

152 useTempRoot = False 

153 

154 

155class ChainedDatastoreButlerTestCase(ButlerFitsTests, lsst.utils.tests.TestCase): 

156 """PosixDatastore specialization""" 

157 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml") 

158 

159 

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

161@mock_s3 

162class S3DatastoreButlerTestCase(ButlerFitsTests, lsst.utils.tests.TestCase): 

163 """S3Datastore specialization of a butler; an S3 storage Datastore + 

164 a local in-memory SqlRegistry. 

165 """ 

166 configFile = os.path.join(TESTDIR, "config/basic/butler-s3store.yaml") 

167 

168 bucketName = "anybucketname" 

169 """Name of the Bucket that will be used in the tests. The name is read from 

170 the config file used with the tests during set-up. 

171 """ 

172 

173 root = "butlerRoot/" 

174 """Root repository directory expected to be used in case useTempRoot=False. 

175 Otherwise the root is set to a 20 characters long randomly generated string 

176 during set-up. 

177 """ 

178 

179 def genRoot(self): 

180 """Returns a random string of len 20 to serve as a root 

181 name for the temporary bucket repo. 

182 

183 This is equivalent to tempfile.mkdtemp as this is what self.root 

184 becomes when useTempRoot is True. 

185 """ 

186 rndstr = "".join( 

187 random.choice(string.ascii_uppercase + string.digits) for _ in range(20) 

188 ) 

189 return rndstr + "/" 

190 

191 def setUp(self): 

192 config = Config(self.configFile) 

193 uri = ButlerURI(config[".datastore.datastore.root"]) 

194 self.bucketName = uri.netloc 

195 

196 if self.useTempRoot: 

197 self.root = self.genRoot() 

198 rooturi = f"s3://{self.bucketName}/{self.root}" 

199 config.update({"datastore": {"datastore": {"root": rooturi}}}) 

200 

201 # set up some fake credentials if they do not exist 

202 self.usingDummyCredentials = setAwsEnvCredentials() 

203 

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

205 # (this used to be the class attribute bucketName) 

206 s3 = boto3.resource("s3") 

207 s3.create_bucket(Bucket=self.bucketName) 

208 

209 self.datastoreStr = f"datastore={self.root}" 

210 self.datastoreName = [f"S3Datastore@{rooturi}"] 

211 Butler.makeRepo(rooturi, config=config, forceConfigRoot=False) 

212 self.tmpConfigFile = os.path.join(rooturi, "butler.yaml") 

213 

214 def tearDown(self): 

215 s3 = boto3.resource("s3") 

216 bucket = s3.Bucket(self.bucketName) 

217 try: 

218 bucket.objects.all().delete() 

219 except botocore.exceptions.ClientError as err: 

220 errorcode = err.response["ResponseMetadata"]["HTTPStatusCode"] 

221 if errorcode == 404: 

222 # the key was not reachable - pass 

223 pass 

224 else: 

225 raise 

226 

227 bucket = s3.Bucket(self.bucketName) 

228 bucket.delete() 

229 

230 # unset any potentially set dummy credentials 

231 if self.usingDummyCredentials: 

232 unsetAwsEnvCredentials() 

233 

234 

235if __name__ == "__main__": 235 ↛ 236line 235 didn't jump to line 236, because the condition on line 235 was never true

236 unittest.main()