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 

22"""Unit tests for daf_butler CLI query-collections command. 

23""" 

24 

25import astropy 

26from astropy.table import Table as AstropyTable 

27from numpy import array 

28import os 

29import shutil 

30import tempfile 

31import unittest 

32 

33from lsst.daf.butler import ( 

34 Butler, 

35 Config, 

36 DatasetRef, 

37 StorageClassFactory 

38) 

39from lsst.daf.butler import script 

40from lsst.daf.butler.tests import addDatasetType, MetricsExample 

41from lsst.daf.butler.tests.utils import ButlerTestHelper 

42 

43 

44TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

45 

46 

47class QueryDataIdsTest(unittest.TestCase, ButlerTestHelper): 

48 

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

50 storageClassFactory = StorageClassFactory() 

51 

52 @staticmethod 

53 def _makeExampleMetrics(): 

54 return MetricsExample({"AM1": 5.2, "AM2": 30.6}, 

55 {"a": [1, 2, 3], 

56 "b": {"blue": 5, "red": "green"}}, 

57 [563, 234, 456.7, 752, 8, 9, 27]) 

58 

59 @staticmethod 

60 def _queryDataIds(repo, dimensions=(), collections=(), datasets=None, where=None): 

61 """Helper to populate the call to script.queryDataIds with default 

62 values.""" 

63 return script.queryDataIds(repo=repo, 

64 dimensions=dimensions, 

65 collections=collections, 

66 datasets=datasets, 

67 where=where) 

68 

69 def setUp(self): 

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

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

72 self.storageClassFactory.addFromConfig(self.configFile) 

73 

74 # New datasets will be added to run and tag, but we will only look in 

75 # tag when looking up datasets. 

76 run = "ingest/run" 

77 tag = "ingest" 

78 self.butler = Butler(butlerConfig, run=run, collections=[tag], tags=[tag]) 

79 

80 # There will not be a collection yet 

81 collections = set(self.butler.registry.queryCollections()) 

82 self.assertEqual(collections, set([run, tag])) 

83 

84 storageClass = self.storageClassFactory.getStorageClass("StructuredCompositeReadComp") 

85 

86 # Create and register a DatasetType 

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

88 datasetTypeName = "test_metric_comp" 

89 self.datasetType = addDatasetType(self.butler, datasetTypeName, dimensions, storageClass) 

90 

91 # Add needed Dimensions 

92 self.butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"}) 

93 self.butler.registry.insertDimensionData("instrument", {"name": "alt_DummyCamComp"}) 

94 self.butler.registry.insertDimensionData("physical_filter", {"instrument": "DummyCamComp", 

95 "name": "d-r", 

96 "band": "R"}) 

97 self.butler.registry.insertDimensionData("visit_system", {"instrument": "DummyCamComp", 

98 "id": 1, 

99 "name": "default"}) 

100 visit_start = astropy.time.Time("2020-01-01 08:00:00.123456789", scale="tai") 

101 visit_end = astropy.time.Time("2020-01-01 08:00:36.66", scale="tai") 

102 self.butler.registry.insertDimensionData("visit", 

103 {"instrument": "DummyCamComp", "id": 423, 

104 "name": "fourtwentythree", "physical_filter": "d-r", 

105 "visit_system": 1, "datetime_begin": visit_start, 

106 "datetime_end": visit_end}) 

107 self.butler.registry.insertDimensionData("visit", {"instrument": "DummyCamComp", "id": 424, 

108 "name": "fourtwentyfour", "physical_filter": "d-r", 

109 "visit_system": 1}) 

110 metric = self._makeExampleMetrics() 

111 dataId = {"instrument": "DummyCamComp", "visit": 423} 

112 ref = DatasetRef(self.datasetType, dataId, id=None) 

113 self.butler.put(metric, ref) 

114 

115 metric = self._makeExampleMetrics() 

116 dataId = {"instrument": "DummyCamComp", "visit": 424} 

117 ref = DatasetRef(self.datasetType, dataId, id=None) 

118 self.butler.put(metric, ref) 

119 

120 def tearDown(self): 

121 if os.path.exists(self.root): 

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

123 

124 def testDimensions(self): 

125 """Test getting a dimension.""" 

126 res = self._queryDataIds(self.root, dimensions=("visit",)) 

127 expected = AstropyTable( 

128 array(( 

129 ("R", "DummyCamComp", "d-r", 1, 423), 

130 ("R", "DummyCamComp", "d-r", 1, 424) 

131 )), 

132 names=("band", "instrument", "physical_filter", "visit_system", "visit") 

133 ) 

134 self.assertAstropyTablesEqual(res, expected) 

135 

136 def testNull(self): 

137 "Test asking for nothing." 

138 res = self._queryDataIds(self.root) 

139 self.assertEqual(res, None) 

140 

141 def testDatasets(self): 

142 """Test getting datasets.""" 

143 res = self._queryDataIds(self.root, datasets="test_metric_comp") 

144 expected = AstropyTable( 

145 array(( 

146 ("R", "DummyCamComp", "d-r", 1, 423), 

147 ("R", "DummyCamComp", "d-r", 1, 424) 

148 )), 

149 names=("band", "instrument", "physical_filter", "visit_system", "visit") 

150 ) 

151 self.assertAstropyTablesEqual(res, expected) 

152 

153 def testWhere(self): 

154 """Test getting datasets.""" 

155 res = self._queryDataIds(self.root, dimensions=("visit",), where="visit=423") 

156 expected = AstropyTable( 

157 array(( 

158 ("R", "DummyCamComp", "d-r", 1, 423), 

159 )), 

160 names=("band", "instrument", "physical_filter", "visit_system", "visit") 

161 ) 

162 self.assertAstropyTablesEqual(res, expected) 

163 

164 def testCollections(self): 

165 """Test getting datasets using the collections option.""" 

166 

167 # Add a dataset in a different collection 

168 butler = Butler(self.root, run="foo") 

169 butler.registry.insertDimensionData("visit", {"instrument": "DummyCamComp", "id": 425, 

170 "name": "fourtwentyfive", "physical_filter": "d-r", 

171 "visit_system": 1}) 

172 metric = self._makeExampleMetrics() 

173 dataId = {"instrument": "DummyCamComp", "visit": 425} 

174 ref = DatasetRef(self.datasetType, dataId, id=None) 

175 butler.put(metric, ref) 

176 

177 # Verify the new dataset is not found in the "ingest/run" collection. 

178 res = self._queryDataIds(repo=self.root, dimensions=("visit",), collections=("ingest/run",), 

179 datasets="test_metric_comp") 

180 expected = AstropyTable( 

181 array(( 

182 ("R", "DummyCamComp", "d-r", 1, 423), 

183 ("R", "DummyCamComp", "d-r", 1, 424) 

184 )), 

185 names=("band", "instrument", "physical_filter", "visit_system", "visit") 

186 ) 

187 self.assertAstropyTablesEqual(res, expected) 

188 

189 # Verify the new dataset is found in the "foo" collection. 

190 res = self._queryDataIds(repo=self.root, dimensions=("visit",), collections=("foo",), 

191 datasets="test_metric_comp") 

192 expected = AstropyTable( 

193 array(( 

194 ("R", "DummyCamComp", "d-r", 1, 425), 

195 )), 

196 names=("band", "instrument", "physical_filter", "visit_system", "visit") 

197 ) 

198 self.assertAstropyTablesEqual(res, expected) 

199 

200 

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

202 unittest.main()