Coverage for python/lsst/dax/apdb/tests/_apdb.py: 20%

Shortcuts 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

165 statements  

1# This file is part of dax_apdb. 

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 

22from __future__ import annotations 

23 

24__all__ = ["ApdbTest"] 

25 

26from abc import ABC, abstractmethod 

27from typing import Any, Optional 

28 

29import pandas 

30 

31from lsst.daf.base import DateTime 

32from lsst.dax.apdb import ApdbConfig, ApdbTables, make_apdb 

33from lsst.sphgeom import Angle, Circle, Region, UnitVector3d 

34from .data_factory import makeObjectCatalog, makeForcedSourceCatalog, makeSourceCatalog 

35 

36 

37class ApdbTest(ABC): 

38 """Base class for Apdb tests that can be specialized for concrete 

39 implementation. 

40 

41 This can only be used as a mixin class for a unittest.TestCase and it 

42 calls various assert methods. 

43 """ 

44 

45 time_partition_tables = False 

46 visit_time = DateTime("2021-01-01T00:00:00", DateTime.TAI) 

47 

48 fsrc_requires_id_list = False 

49 """Should be set to True if getDiaForcedSources requires object IDs""" 

50 

51 # number of columns as defined in schema YAML files 

52 n_obj_columns = 91 + 2 # schema + schema-extra 

53 n_obj_last_columns = 17 

54 n_src_columns = 107 

55 n_fsrc_columns = 8 

56 

57 @abstractmethod 

58 def make_config(self, **kwargs: Any) -> ApdbConfig: 

59 """Make config class instance used in all tests.""" 

60 raise NotImplementedError() 

61 

62 @abstractmethod 

63 def n_columns(self, table: ApdbTables) -> int: 

64 """Return number of columns for a specified table.""" 

65 raise NotImplementedError() 

66 

67 def make_region(self) -> Region: 

68 """Make a region used in tests""" 

69 pointing_v = UnitVector3d(1., 1., -1.) 

70 fov = 0.05 # radians 

71 region = Circle(pointing_v, Angle(fov/2)) 

72 return region 

73 

74 def assert_catalog(self, catalog: Any, rows: int, table: ApdbTables) -> None: 

75 """Validate catalog type and size 

76 

77 Parameters 

78 ---------- 

79 catalog : `object` 

80 Expected type of this is ``type``. 

81 rows : int 

82 Expected number of rows in a catalog. 

83 table : `ApdbTables` 

84 APDB table type. 

85 """ 

86 self.assertIsInstance(catalog, pandas.DataFrame) # type: ignore[attr-defined] 

87 self.assertEqual(catalog.shape[0], rows) # type: ignore[attr-defined] 

88 self.assertEqual(catalog.shape[1], self.n_columns(table)) # type: ignore[attr-defined] 

89 

90 def test_makeSchema(self) -> None: 

91 """Test for makeing APDB schema.""" 

92 config = self.make_config() 

93 apdb = make_apdb(config) 

94 

95 apdb.makeSchema() 

96 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject)) # type: ignore[attr-defined] 

97 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObjectLast)) # type: ignore[attr-defined] 

98 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaSource)) # type: ignore[attr-defined] 

99 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaForcedSource)) # type: ignore[attr-defined] 

100 

101 def test_empty_gets(self) -> None: 

102 """Test for getting data from empty database. 

103 

104 All get() methods should return empty results, only useful for 

105 checking that code is not broken. 

106 """ 

107 

108 # use non-zero months for Forced/Source fetching 

109 config = self.make_config() 

110 apdb = make_apdb(config) 

111 apdb.makeSchema() 

112 

113 region = self.make_region() 

114 visit_time = self.visit_time 

115 

116 res: Optional[pandas.DataFrame] 

117 

118 # get objects by region 

119 res = apdb.getDiaObjects(region) 

120 self.assert_catalog(res, 0, ApdbTables.DiaObject) 

121 

122 # get sources by region 

123 res = apdb.getDiaSources(region, None, visit_time) 

124 self.assert_catalog(res, 0, ApdbTables.DiaSource) 

125 

126 res = apdb.getDiaSources(region, [], visit_time) 

127 self.assert_catalog(res, 0, ApdbTables.DiaSource) 

128 

129 # get sources by object ID, non-empty object list 

130 res = apdb.getDiaSources(region, [1, 2, 3], visit_time) 

131 self.assert_catalog(res, 0, ApdbTables.DiaSource) 

132 

133 # get forced sources by object ID, empty object list 

134 res = apdb.getDiaForcedSources(region, [], visit_time) 

135 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource) 

136 

137 # get sources by object ID, non-empty object list 

138 res = apdb.getDiaForcedSources(region, [1, 2, 3], visit_time) 

139 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource) 

140 

141 # get sources by region 

142 if self.fsrc_requires_id_list: 

143 with self.assertRaises(NotImplementedError): # type: ignore[attr-defined] 

144 apdb.getDiaForcedSources(region, None, visit_time) 

145 else: 

146 apdb.getDiaForcedSources(region, None, visit_time) 

147 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource) 

148 

149 def test_empty_gets_0months(self) -> None: 

150 """Test for getting data from empty database. 

151 

152 All get() methods should return empty DataFrame or None. 

153 """ 

154 

155 # set read_sources_months to 0 so that Forced/Sources are None 

156 config = self.make_config(read_sources_months=0, 

157 read_forced_sources_months=0) 

158 apdb = make_apdb(config) 

159 apdb.makeSchema() 

160 

161 region = self.make_region() 

162 visit_time = self.visit_time 

163 

164 res: Optional[pandas.DataFrame] 

165 

166 # get objects by region 

167 res = apdb.getDiaObjects(region) 

168 self.assert_catalog(res, 0, ApdbTables.DiaObject) 

169 

170 # get sources by region 

171 res = apdb.getDiaSources(region, None, visit_time) 

172 self.assertIs(res, None) # type: ignore[attr-defined] 

173 

174 # get sources by object ID, empty object list 

175 res = apdb.getDiaSources(region, [], visit_time) 

176 self.assertIs(res, None) # type: ignore[attr-defined] 

177 

178 # get forced sources by object ID, empty object list 

179 res = apdb.getDiaForcedSources(region, [], visit_time) 

180 self.assertIs(res, None) # type: ignore[attr-defined] 

181 

182 def test_storeObjects(self) -> None: 

183 """Store and retrieve DiaObjects.""" 

184 

185 # don't care about sources. 

186 config = self.make_config() 

187 apdb = make_apdb(config) 

188 apdb.makeSchema() 

189 

190 region = self.make_region() 

191 visit_time = self.visit_time 

192 

193 # make catalog with Objects 

194 catalog = makeObjectCatalog(region, 100) 

195 

196 # store catalog 

197 apdb.store(visit_time, catalog) 

198 

199 # read it back and check sizes 

200 res = apdb.getDiaObjects(region) 

201 self.assert_catalog(res, len(catalog), ApdbTables.DiaObject) 

202 

203 def test_storeSources(self) -> None: 

204 """Store and retrieve DiaSources.""" 

205 config = self.make_config() 

206 apdb = make_apdb(config) 

207 apdb.makeSchema() 

208 

209 region = self.make_region() 

210 visit_time = self.visit_time 

211 

212 # have to store Objects first 

213 objects = makeObjectCatalog(region, 100) 

214 oids = list(objects["diaObjectId"]) 

215 sources = makeSourceCatalog(objects, visit_time) 

216 

217 # save the objects and sources 

218 apdb.store(visit_time, objects, sources) 

219 

220 # read it back, no ID filtering 

221 res = apdb.getDiaSources(region, None, visit_time) 

222 self.assert_catalog(res, len(sources), ApdbTables.DiaSource) 

223 

224 # read it back and filter by ID 

225 res = apdb.getDiaSources(region, oids, visit_time) 

226 self.assert_catalog(res, len(sources), ApdbTables.DiaSource) 

227 

228 # read it back to get schema 

229 res = apdb.getDiaSources(region, [], visit_time) 

230 self.assert_catalog(res, 0, ApdbTables.DiaSource) 

231 

232 def test_storeForcedSources(self) -> None: 

233 """Store and retrieve DiaForcedSources.""" 

234 

235 config = self.make_config() 

236 apdb = make_apdb(config) 

237 apdb.makeSchema() 

238 

239 region = self.make_region() 

240 visit_time = self.visit_time 

241 

242 # have to store Objects first 

243 objects = makeObjectCatalog(region, 100) 

244 oids = list(objects["diaObjectId"]) 

245 catalog = makeForcedSourceCatalog(objects, visit_time) 

246 

247 apdb.store(visit_time, objects, forced_sources=catalog) 

248 

249 # read it back and check sizes 

250 res = apdb.getDiaForcedSources(region, oids, visit_time) 

251 self.assert_catalog(res, len(catalog), ApdbTables.DiaForcedSource) 

252 

253 # read it back to get schema 

254 res = apdb.getDiaForcedSources(region, [], visit_time) 

255 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource) 

256 

257 def test_midPointTai_src(self) -> None: 

258 """Test for time filtering of DiaSources. 

259 """ 

260 config = self.make_config() 

261 apdb = make_apdb(config) 

262 apdb.makeSchema() 

263 

264 region = self.make_region() 

265 # 2021-01-01 plus 360 days is 2021-12-27 

266 src_time1 = DateTime("2021-01-01T00:00:00", DateTime.TAI) 

267 src_time2 = DateTime("2021-01-01T00:00:02", DateTime.TAI) 

268 visit_time0 = DateTime("2021-12-26T23:59:59", DateTime.TAI) 

269 visit_time1 = DateTime("2021-12-27T00:00:01", DateTime.TAI) 

270 visit_time2 = DateTime("2021-12-27T00:00:03", DateTime.TAI) 

271 

272 objects = makeObjectCatalog(region, 100) 

273 oids = list(objects["diaObjectId"]) 

274 sources = makeSourceCatalog(objects, src_time1, 0) 

275 apdb.store(src_time1, objects, sources) 

276 

277 sources = makeSourceCatalog(objects, src_time2, 100) 

278 apdb.store(src_time2, objects, sources) 

279 

280 # reading at time of last save should read all 

281 res = apdb.getDiaSources(region, oids, src_time2) 

282 self.assert_catalog(res, 200, ApdbTables.DiaSource) 

283 

284 # one second before 12 months 

285 res = apdb.getDiaSources(region, oids, visit_time0) 

286 self.assert_catalog(res, 200, ApdbTables.DiaSource) 

287 

288 # reading at later time of last save should only read a subset 

289 res = apdb.getDiaSources(region, oids, visit_time1) 

290 self.assert_catalog(res, 100, ApdbTables.DiaSource) 

291 

292 # reading at later time of last save should only read a subset 

293 res = apdb.getDiaSources(region, oids, visit_time2) 

294 self.assert_catalog(res, 0, ApdbTables.DiaSource) 

295 

296 def test_midPointTai_fsrc(self) -> None: 

297 """Test for time filtering of DiaForcedSources. 

298 """ 

299 config = self.make_config() 

300 apdb = make_apdb(config) 

301 apdb.makeSchema() 

302 

303 region = self.make_region() 

304 src_time1 = DateTime("2021-01-01T00:00:00", DateTime.TAI) 

305 src_time2 = DateTime("2021-01-01T00:00:02", DateTime.TAI) 

306 visit_time0 = DateTime("2021-12-26T23:59:59", DateTime.TAI) 

307 visit_time1 = DateTime("2021-12-27T00:00:01", DateTime.TAI) 

308 visit_time2 = DateTime("2021-12-27T00:00:03", DateTime.TAI) 

309 

310 objects = makeObjectCatalog(region, 100) 

311 oids = list(objects["diaObjectId"]) 

312 sources = makeForcedSourceCatalog(objects, src_time1, 1) 

313 apdb.store(src_time1, objects, forced_sources=sources) 

314 

315 sources = makeForcedSourceCatalog(objects, src_time2, 2) 

316 apdb.store(src_time2, objects, forced_sources=sources) 

317 

318 # reading at time of last save should read all 

319 res = apdb.getDiaForcedSources(region, oids, src_time2) 

320 self.assert_catalog(res, 200, ApdbTables.DiaForcedSource) 

321 

322 # one second before 12 months 

323 res = apdb.getDiaForcedSources(region, oids, visit_time0) 

324 self.assert_catalog(res, 200, ApdbTables.DiaForcedSource) 

325 

326 # reading at later time of last save should only read a subset 

327 res = apdb.getDiaForcedSources(region, oids, visit_time1) 

328 self.assert_catalog(res, 100, ApdbTables.DiaForcedSource) 

329 

330 # reading at later time of last save should only read a subset 

331 res = apdb.getDiaForcedSources(region, oids, visit_time2) 

332 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource)