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

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

336 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, Tuple 

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, makeSSObjectCatalog 

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 fsrc_history_region_filtering = False 

52 """Should be set to True if forced sources history support region-based 

53 filtering. 

54 """ 

55 

56 # number of columns as defined in schema YAML files 

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

58 n_obj_last_columns = 17 

59 n_src_columns = 107 

60 n_fsrc_columns = 8 

61 n_ssobj_columns = 81 

62 

63 @abstractmethod 

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

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

66 raise NotImplementedError() 

67 

68 @abstractmethod 

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

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

71 raise NotImplementedError() 

72 

73 @abstractmethod 

74 def getDiaObjects_table(self) -> ApdbTables: 

75 """Return type of table returned from getDiaObjects method.""" 

76 raise NotImplementedError() 

77 

78 def make_region(self, xyz: Tuple[float, float, float] = (1., 1., -1.)) -> Region: 

79 """Make a region to use in tests""" 

80 pointing_v = UnitVector3d(*xyz) 

81 fov = 0.05 # radians 

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

83 return region 

84 

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

86 """Validate catalog type and size 

87 

88 Parameters 

89 ---------- 

90 catalog : `object` 

91 Expected type of this is ``type``. 

92 rows : int 

93 Expected number of rows in a catalog. 

94 table : `ApdbTables` 

95 APDB table type. 

96 """ 

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

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

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

100 

101 def test_makeSchema(self) -> None: 

102 """Test for makeing APDB schema.""" 

103 config = self.make_config() 

104 apdb = make_apdb(config) 

105 

106 apdb.makeSchema() 

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

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

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

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

111 

112 def test_empty_gets(self) -> None: 

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

114 

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

116 checking that code is not broken. 

117 """ 

118 

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

120 config = self.make_config() 

121 apdb = make_apdb(config) 

122 apdb.makeSchema() 

123 

124 region = self.make_region() 

125 visit_time = self.visit_time 

126 

127 res: Optional[pandas.DataFrame] 

128 

129 # get objects by region 

130 res = apdb.getDiaObjects(region) 

131 self.assert_catalog(res, 0, self.getDiaObjects_table()) 

132 

133 # get sources by region 

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

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

136 

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

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

139 

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

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

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

143 

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

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

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

147 

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

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

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

151 

152 # get sources by region 

153 if self.fsrc_requires_id_list: 

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

155 apdb.getDiaForcedSources(region, None, visit_time) 

156 else: 

157 apdb.getDiaForcedSources(region, None, visit_time) 

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

159 

160 def test_empty_gets_0months(self) -> None: 

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

162 

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

164 """ 

165 

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

167 config = self.make_config(read_sources_months=0, 

168 read_forced_sources_months=0) 

169 apdb = make_apdb(config) 

170 apdb.makeSchema() 

171 

172 region = self.make_region() 

173 visit_time = self.visit_time 

174 

175 res: Optional[pandas.DataFrame] 

176 

177 # get objects by region 

178 res = apdb.getDiaObjects(region) 

179 self.assert_catalog(res, 0, self.getDiaObjects_table()) 

180 

181 # get sources by region 

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

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

184 

185 # get sources by object ID, empty object list 

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

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

188 

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

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

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

192 

193 def test_storeObjects(self) -> None: 

194 """Store and retrieve DiaObjects.""" 

195 

196 # don't care about sources. 

197 config = self.make_config() 

198 apdb = make_apdb(config) 

199 apdb.makeSchema() 

200 

201 region = self.make_region() 

202 visit_time = self.visit_time 

203 

204 # make catalog with Objects 

205 catalog = makeObjectCatalog(region, 100) 

206 

207 # store catalog 

208 apdb.store(visit_time, catalog) 

209 

210 # read it back and check sizes 

211 res = apdb.getDiaObjects(region) 

212 self.assert_catalog(res, len(catalog), self.getDiaObjects_table()) 

213 

214 def test_objectHistory(self) -> None: 

215 """Store and retrieve DiaObject history.""" 

216 

217 # don't care about sources. 

218 config = self.make_config() 

219 apdb = make_apdb(config) 

220 apdb.makeSchema() 

221 

222 region1 = self.make_region((1., 1., -1.)) 

223 region2 = self.make_region((-1., -1., -1.)) 

224 visit_time = [ 

225 DateTime("2021-01-01T00:01:00", DateTime.TAI), 

226 DateTime("2021-01-01T00:02:00", DateTime.TAI), 

227 DateTime("2021-01-01T00:03:00", DateTime.TAI), 

228 DateTime("2021-01-01T00:04:00", DateTime.TAI), 

229 DateTime("2021-01-01T00:05:00", DateTime.TAI), 

230 DateTime("2021-01-01T00:06:00", DateTime.TAI), 

231 DateTime("2021-03-01T00:01:00", DateTime.TAI), 

232 DateTime("2021-03-01T00:02:00", DateTime.TAI), 

233 ] 

234 end_time = DateTime("2021-03-02T00:00:00", DateTime.TAI) 

235 

236 nobj = 100 

237 catalog1 = makeObjectCatalog(region1, nobj) 

238 apdb.store(visit_time[0], catalog1) 

239 apdb.store(visit_time[2], catalog1) 

240 apdb.store(visit_time[4], catalog1) 

241 apdb.store(visit_time[6], catalog1) 

242 catalog2 = makeObjectCatalog(region2, nobj, start_id=nobj*2) 

243 apdb.store(visit_time[1], catalog2) 

244 apdb.store(visit_time[3], catalog2) 

245 apdb.store(visit_time[5], catalog2) 

246 apdb.store(visit_time[7], catalog2) 

247 

248 # read it back and check sizes 

249 res = apdb.getDiaObjectsHistory(DateTime("2021-01-01T00:00:00", DateTime.TAI), end_time) 

250 self.assert_catalog(res, nobj * 8, ApdbTables.DiaObject) 

251 

252 res = apdb.getDiaObjectsHistory(DateTime("2021-01-01T00:01:00", DateTime.TAI), end_time) 

253 self.assert_catalog(res, nobj * 8, ApdbTables.DiaObject) 

254 

255 res = apdb.getDiaObjectsHistory(DateTime("2021-01-01T00:01:01", DateTime.TAI), end_time) 

256 self.assert_catalog(res, nobj * 7, ApdbTables.DiaObject) 

257 

258 res = apdb.getDiaObjectsHistory(DateTime("2021-01-01T00:02:30", DateTime.TAI), end_time) 

259 self.assert_catalog(res, nobj * 6, ApdbTables.DiaObject) 

260 

261 res = apdb.getDiaObjectsHistory(DateTime("2021-01-01T00:05:00", DateTime.TAI), end_time) 

262 self.assert_catalog(res, nobj * 4, ApdbTables.DiaObject) 

263 

264 res = apdb.getDiaObjectsHistory(DateTime("2021-01-01T00:06:30", DateTime.TAI), end_time) 

265 self.assert_catalog(res, nobj * 2, ApdbTables.DiaObject) 

266 

267 res = apdb.getDiaObjectsHistory(DateTime("2021-03-01T00:02:00.001", DateTime.TAI), end_time) 

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

269 

270 res = apdb.getDiaObjectsHistory( 

271 DateTime("2021-01-01T00:00:00", DateTime.TAI), 

272 DateTime("2021-01-01T00:06:01", DateTime.TAI), 

273 ) 

274 self.assert_catalog(res, nobj * 6, ApdbTables.DiaObject) 

275 

276 res = apdb.getDiaObjectsHistory( 

277 DateTime("2021-01-01T00:00:00", DateTime.TAI), 

278 DateTime("2021-01-01T00:06:00", DateTime.TAI), 

279 ) 

280 self.assert_catalog(res, nobj * 5, ApdbTables.DiaObject) 

281 

282 res = apdb.getDiaObjectsHistory( 

283 DateTime("2021-01-01T00:00:00", DateTime.TAI), 

284 DateTime("2021-01-01T00:01:00", DateTime.TAI), 

285 ) 

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

287 

288 res = apdb.getDiaObjectsHistory( 

289 DateTime("2021-01-01T00:00:00", DateTime.TAI), end_time, region=region1 

290 ) 

291 self.assert_catalog(res, nobj * 4, ApdbTables.DiaObject) 

292 

293 res = apdb.getDiaObjectsHistory( 

294 DateTime("2021-01-01T00:03:00", DateTime.TAI), end_time, region=region2 

295 ) 

296 self.assert_catalog(res, nobj * 3, ApdbTables.DiaObject) 

297 

298 res = apdb.getDiaObjectsHistory( 

299 DateTime("2021-01-01T00:00:00", DateTime.TAI), 

300 DateTime("2021-01-01T00:03:30", DateTime.TAI), 

301 region1, 

302 ) 

303 self.assert_catalog(res, nobj * 2, ApdbTables.DiaObject) 

304 

305 def test_storeSources(self) -> None: 

306 """Store and retrieve DiaSources.""" 

307 config = self.make_config() 

308 apdb = make_apdb(config) 

309 apdb.makeSchema() 

310 

311 region = self.make_region() 

312 visit_time = self.visit_time 

313 

314 # have to store Objects first 

315 objects = makeObjectCatalog(region, 100) 

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

317 sources = makeSourceCatalog(objects, visit_time) 

318 

319 # save the objects and sources 

320 apdb.store(visit_time, objects, sources) 

321 

322 # read it back, no ID filtering 

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

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

325 

326 # read it back and filter by ID 

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

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

329 

330 # read it back to get schema 

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

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

333 

334 def test_sourceHistory(self) -> None: 

335 """Store and retrieve DiaSource history.""" 

336 

337 # don't care about sources. 

338 config = self.make_config() 

339 apdb = make_apdb(config) 

340 apdb.makeSchema() 

341 

342 region1 = self.make_region((1., 1., -1.)) 

343 region2 = self.make_region((-1., -1., -1.)) 

344 nobj = 100 

345 objects1 = makeObjectCatalog(region1, nobj) 

346 objects2 = makeObjectCatalog(region2, nobj, start_id=nobj*2) 

347 

348 visits = [ 

349 (DateTime("2021-01-01T00:01:00", DateTime.TAI), objects1), 

350 (DateTime("2021-01-01T00:02:00", DateTime.TAI), objects2), 

351 (DateTime("2021-01-01T00:03:00", DateTime.TAI), objects1), 

352 (DateTime("2021-01-01T00:04:00", DateTime.TAI), objects2), 

353 (DateTime("2021-01-01T00:05:00", DateTime.TAI), objects1), 

354 (DateTime("2021-01-01T00:06:00", DateTime.TAI), objects2), 

355 (DateTime("2021-03-01T00:01:00", DateTime.TAI), objects1), 

356 (DateTime("2021-03-01T00:02:00", DateTime.TAI), objects2), 

357 ] 

358 end_time = DateTime("2021-03-02T00:00:00", DateTime.TAI) 

359 

360 start_id = 0 

361 for visit_time, objects in visits: 

362 sources = makeSourceCatalog(objects, visit_time, start_id=start_id) 

363 apdb.store(visit_time, objects, sources) 

364 start_id += nobj 

365 

366 # read it back and check sizes 

367 res = apdb.getDiaSourcesHistory(DateTime("2021-01-01T00:00:00", DateTime.TAI), end_time) 

368 self.assert_catalog(res, nobj * 8, ApdbTables.DiaSource) 

369 

370 res = apdb.getDiaSourcesHistory(DateTime("2021-01-01T00:01:00", DateTime.TAI), end_time) 

371 self.assert_catalog(res, nobj * 8, ApdbTables.DiaSource) 

372 

373 res = apdb.getDiaSourcesHistory(DateTime("2021-01-01T00:01:01", DateTime.TAI), end_time) 

374 self.assert_catalog(res, nobj * 7, ApdbTables.DiaSource) 

375 

376 res = apdb.getDiaSourcesHistory(DateTime("2021-01-01T00:02:30", DateTime.TAI), end_time) 

377 self.assert_catalog(res, nobj * 6, ApdbTables.DiaSource) 

378 

379 res = apdb.getDiaSourcesHistory(DateTime("2021-01-01T00:05:00", DateTime.TAI), end_time) 

380 self.assert_catalog(res, nobj * 4, ApdbTables.DiaSource) 

381 

382 res = apdb.getDiaSourcesHistory(DateTime("2021-01-01T00:06:30", DateTime.TAI), end_time) 

383 self.assert_catalog(res, nobj * 2, ApdbTables.DiaSource) 

384 

385 res = apdb.getDiaSourcesHistory(DateTime("2021-03-01T00:02:00.001", DateTime.TAI), end_time) 

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

387 

388 res = apdb.getDiaSourcesHistory( 

389 DateTime("2021-01-01T00:00:00", DateTime.TAI), 

390 DateTime("2021-01-01T00:06:01", DateTime.TAI), 

391 ) 

392 self.assert_catalog(res, nobj * 6, ApdbTables.DiaSource) 

393 

394 res = apdb.getDiaSourcesHistory( 

395 DateTime("2021-01-01T00:00:00", DateTime.TAI), 

396 DateTime("2021-01-01T00:06:00", DateTime.TAI), 

397 ) 

398 self.assert_catalog(res, nobj * 5, ApdbTables.DiaSource) 

399 

400 res = apdb.getDiaSourcesHistory( 

401 DateTime("2021-01-01T00:00:00", DateTime.TAI), 

402 DateTime("2021-01-01T00:01:00", DateTime.TAI), 

403 ) 

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

405 

406 res = apdb.getDiaSourcesHistory( 

407 DateTime("2021-01-01T00:00:00", DateTime.TAI), end_time, region=region1 

408 ) 

409 self.assert_catalog(res, nobj * 4, ApdbTables.DiaSource) 

410 

411 res = apdb.getDiaSourcesHistory( 

412 DateTime("2021-01-01T00:03:00", DateTime.TAI), end_time, region=region2 

413 ) 

414 self.assert_catalog(res, nobj * 3, ApdbTables.DiaSource) 

415 

416 res = apdb.getDiaSourcesHistory( 

417 DateTime("2021-01-01T00:00:00", DateTime.TAI), 

418 DateTime("2021-01-01T00:03:30", DateTime.TAI), 

419 region1, 

420 ) 

421 self.assert_catalog(res, nobj * 2, ApdbTables.DiaSource) 

422 

423 def test_storeForcedSources(self) -> None: 

424 """Store and retrieve DiaForcedSources.""" 

425 

426 config = self.make_config() 

427 apdb = make_apdb(config) 

428 apdb.makeSchema() 

429 

430 region = self.make_region() 

431 visit_time = self.visit_time 

432 

433 # have to store Objects first 

434 objects = makeObjectCatalog(region, 100) 

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

436 catalog = makeForcedSourceCatalog(objects, visit_time) 

437 

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

439 

440 # read it back and check sizes 

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

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

443 

444 # read it back to get schema 

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

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

447 

448 def test_forcedSourceHistory(self) -> None: 

449 """Store and retrieve DiaForcedSource history.""" 

450 

451 # don't care about sources. 

452 config = self.make_config() 

453 apdb = make_apdb(config) 

454 apdb.makeSchema() 

455 

456 region1 = self.make_region((1., 1., -1.)) 

457 region2 = self.make_region((-1., -1., -1.)) 

458 nobj = 100 

459 objects1 = makeObjectCatalog(region1, nobj) 

460 objects2 = makeObjectCatalog(region2, nobj, start_id=nobj*2) 

461 

462 visits = [ 

463 (DateTime("2021-01-01T00:01:00", DateTime.TAI), objects1), 

464 (DateTime("2021-01-01T00:02:00", DateTime.TAI), objects2), 

465 (DateTime("2021-01-01T00:03:00", DateTime.TAI), objects1), 

466 (DateTime("2021-01-01T00:04:00", DateTime.TAI), objects2), 

467 (DateTime("2021-01-01T00:05:00", DateTime.TAI), objects1), 

468 (DateTime("2021-01-01T00:06:00", DateTime.TAI), objects2), 

469 (DateTime("2021-03-01T00:01:00", DateTime.TAI), objects1), 

470 (DateTime("2021-03-01T00:02:00", DateTime.TAI), objects2), 

471 ] 

472 end_time = DateTime("2021-03-02T00:00:00", DateTime.TAI) 

473 

474 start_id = 0 

475 for visit_time, objects in visits: 

476 sources = makeForcedSourceCatalog(objects, visit_time, ccdVisitId=start_id) 

477 apdb.store(visit_time, objects, forced_sources=sources) 

478 start_id += 1 

479 

480 # read it back and check sizes 

481 res = apdb.getDiaForcedSourcesHistory(DateTime("2021-01-01T00:00:00", DateTime.TAI), end_time) 

482 self.assert_catalog(res, nobj * 8, ApdbTables.DiaForcedSource) 

483 

484 res = apdb.getDiaForcedSourcesHistory(DateTime("2021-01-01T00:01:00", DateTime.TAI), end_time) 

485 self.assert_catalog(res, nobj * 8, ApdbTables.DiaForcedSource) 

486 

487 res = apdb.getDiaForcedSourcesHistory(DateTime("2021-01-01T00:01:01", DateTime.TAI), end_time) 

488 self.assert_catalog(res, nobj * 7, ApdbTables.DiaForcedSource) 

489 

490 res = apdb.getDiaForcedSourcesHistory(DateTime("2021-01-01T00:02:30", DateTime.TAI), end_time) 

491 self.assert_catalog(res, nobj * 6, ApdbTables.DiaForcedSource) 

492 

493 res = apdb.getDiaForcedSourcesHistory(DateTime("2021-01-01T00:05:00", DateTime.TAI), end_time) 

494 self.assert_catalog(res, nobj * 4, ApdbTables.DiaForcedSource) 

495 

496 res = apdb.getDiaForcedSourcesHistory(DateTime("2021-01-01T00:06:30", DateTime.TAI), end_time) 

497 self.assert_catalog(res, nobj * 2, ApdbTables.DiaForcedSource) 

498 

499 res = apdb.getDiaForcedSourcesHistory(DateTime("2021-03-01T00:02:00.001", DateTime.TAI), end_time) 

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

501 

502 res = apdb.getDiaForcedSourcesHistory( 

503 DateTime("2021-01-01T00:00:00", DateTime.TAI), 

504 DateTime("2021-01-01T00:06:01", DateTime.TAI), 

505 ) 

506 self.assert_catalog(res, nobj * 6, ApdbTables.DiaForcedSource) 

507 

508 res = apdb.getDiaForcedSourcesHistory( 

509 DateTime("2021-01-01T00:00:00", DateTime.TAI), 

510 DateTime("2021-01-01T00:06:00", DateTime.TAI), 

511 ) 

512 self.assert_catalog(res, nobj * 5, ApdbTables.DiaForcedSource) 

513 

514 res = apdb.getDiaForcedSourcesHistory( 

515 DateTime("2021-01-01T00:00:00", DateTime.TAI), 

516 DateTime("2021-01-01T00:01:00", DateTime.TAI), 

517 ) 

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

519 

520 res = apdb.getDiaForcedSourcesHistory( 

521 DateTime("2021-01-01T00:00:00", DateTime.TAI), end_time, region=region1 

522 ) 

523 rows = nobj * 4 if self.fsrc_history_region_filtering else nobj * 8 

524 self.assert_catalog(res, rows, ApdbTables.DiaForcedSource) 

525 

526 res = apdb.getDiaForcedSourcesHistory( 

527 DateTime("2021-01-01T00:03:00", DateTime.TAI), end_time, region=region2 

528 ) 

529 rows = nobj * 3 if self.fsrc_history_region_filtering else nobj * 6 

530 self.assert_catalog(res, rows, ApdbTables.DiaForcedSource) 

531 

532 res = apdb.getDiaForcedSourcesHistory( 

533 DateTime("2021-01-01T00:00:00", DateTime.TAI), 

534 DateTime("2021-01-01T00:03:30", DateTime.TAI), 

535 region1, 

536 ) 

537 rows = nobj * 2 if self.fsrc_history_region_filtering else nobj * 3 

538 self.assert_catalog(res, rows, ApdbTables.DiaForcedSource) 

539 

540 def test_storeSSObjects(self) -> None: 

541 """Store and retrieve SSObjects.""" 

542 

543 # don't care about sources. 

544 config = self.make_config() 

545 apdb = make_apdb(config) 

546 apdb.makeSchema() 

547 

548 # make catalog with SSObjects 

549 catalog = makeSSObjectCatalog(100, flags=1) 

550 

551 # store catalog 

552 apdb.storeSSObjects(catalog) 

553 

554 # read it back and check sizes 

555 res = apdb.getSSObjects() 

556 self.assert_catalog(res, len(catalog), ApdbTables.SSObject) 

557 

558 # check that override works, make catalog with SSObjects, ID = 51-150 

559 catalog = makeSSObjectCatalog(100, 51, flags=2) 

560 apdb.storeSSObjects(catalog) 

561 res = apdb.getSSObjects() 

562 self.assert_catalog(res, 150, ApdbTables.SSObject) 

563 self.assertEqual(len(res[res["flags"] == 1]), 50) # type: ignore[attr-defined] 

564 self.assertEqual(len(res[res["flags"] == 2]), 100) # type: ignore[attr-defined] 

565 

566 def test_reassignObjects(self) -> None: 

567 """Reassign DiaObjects.""" 

568 

569 # don't care about sources. 

570 config = self.make_config() 

571 apdb = make_apdb(config) 

572 apdb.makeSchema() 

573 

574 region = self.make_region() 

575 visit_time = self.visit_time 

576 objects = makeObjectCatalog(region, 100) 

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

578 sources = makeSourceCatalog(objects, visit_time) 

579 apdb.store(visit_time, objects, sources) 

580 

581 catalog = makeSSObjectCatalog(100) 

582 apdb.storeSSObjects(catalog) 

583 

584 # read it back and filter by ID 

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

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

587 

588 apdb.reassignDiaSources({1: 1, 2: 2, 5: 5}) 

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

590 self.assert_catalog(res, len(sources) - 3, ApdbTables.DiaSource) 

591 

592 with self.assertRaisesRegex(ValueError, r"do not exist.*\D1000"): # type: ignore[attr-defined] 

593 apdb.reassignDiaSources({1000: 1, 7: 3, }) 

594 self.assert_catalog(res, len(sources) - 3, ApdbTables.DiaSource) 

595 

596 def test_midPointTai_src(self) -> None: 

597 """Test for time filtering of DiaSources. 

598 """ 

599 config = self.make_config() 

600 apdb = make_apdb(config) 

601 apdb.makeSchema() 

602 

603 region = self.make_region() 

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

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

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

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

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

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

610 

611 objects = makeObjectCatalog(region, 100) 

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

613 sources = makeSourceCatalog(objects, src_time1, 0) 

614 apdb.store(src_time1, objects, sources) 

615 

616 sources = makeSourceCatalog(objects, src_time2, 100) 

617 apdb.store(src_time2, objects, sources) 

618 

619 # reading at time of last save should read all 

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

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

622 

623 # one second before 12 months 

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

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

626 

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

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

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

630 

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

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

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

634 

635 def test_midPointTai_fsrc(self) -> None: 

636 """Test for time filtering of DiaForcedSources. 

637 """ 

638 config = self.make_config() 

639 apdb = make_apdb(config) 

640 apdb.makeSchema() 

641 

642 region = self.make_region() 

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

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

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

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

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

648 

649 objects = makeObjectCatalog(region, 100) 

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

651 sources = makeForcedSourceCatalog(objects, src_time1, 1) 

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

653 

654 sources = makeForcedSourceCatalog(objects, src_time2, 2) 

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

656 

657 # reading at time of last save should read all 

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

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

660 

661 # one second before 12 months 

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

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

664 

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

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

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

668 

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

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

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