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

342 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-12-30 01:30 -0800

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, Dict, Optional, Tuple 

28 

29import pandas 

30from lsst.daf.base import DateTime 

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

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

33 

34from .data_factory import makeForcedSourceCatalog, makeObjectCatalog, 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 extra_object_columns: Dict[str, Any] = {} 

57 """Additional columns with values to add to DiaObject catalog.""" 

58 

59 # number of columns as defined in tests/config/schema.yaml 

60 n_obj_columns = 8 

61 n_obj_last_columns = 5 

62 n_src_columns = 10 

63 n_fsrc_columns = 4 

64 n_ssobj_columns = 3 

65 

66 @abstractmethod 

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

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

69 raise NotImplementedError() 

70 

71 @abstractmethod 

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

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

74 raise NotImplementedError() 

75 

76 @abstractmethod 

77 def getDiaObjects_table(self) -> ApdbTables: 

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

79 raise NotImplementedError() 

80 

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

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

83 pointing_v = UnitVector3d(*xyz) 

84 fov = 0.05 # radians 

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

86 return region 

87 

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

89 """Validate catalog type and size 

90 

91 Parameters 

92 ---------- 

93 catalog : `object` 

94 Expected type of this is ``type``. 

95 rows : int 

96 Expected number of rows in a catalog. 

97 table : `ApdbTables` 

98 APDB table type. 

99 """ 

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

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

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

103 

104 def test_makeSchema(self) -> None: 

105 """Test for makeing APDB schema.""" 

106 config = self.make_config() 

107 apdb = make_apdb(config) 

108 

109 apdb.makeSchema() 

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

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

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

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

114 

115 def test_empty_gets(self) -> None: 

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

117 

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

119 checking that code is not broken. 

120 """ 

121 

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

123 config = self.make_config() 

124 apdb = make_apdb(config) 

125 apdb.makeSchema() 

126 

127 region = self.make_region() 

128 visit_time = self.visit_time 

129 

130 res: Optional[pandas.DataFrame] 

131 

132 # get objects by region 

133 res = apdb.getDiaObjects(region) 

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

135 

136 # get sources by region 

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

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

139 

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

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

142 

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

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

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

146 

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

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

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

150 

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

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

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

154 

155 # get sources by region 

156 if self.fsrc_requires_id_list: 

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

158 apdb.getDiaForcedSources(region, None, visit_time) 

159 else: 

160 apdb.getDiaForcedSources(region, None, visit_time) 

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

162 

163 def test_empty_gets_0months(self) -> None: 

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

165 

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

167 """ 

168 

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

170 config = self.make_config(read_sources_months=0, 

171 read_forced_sources_months=0) 

172 apdb = make_apdb(config) 

173 apdb.makeSchema() 

174 

175 region = self.make_region() 

176 visit_time = self.visit_time 

177 

178 res: Optional[pandas.DataFrame] 

179 

180 # get objects by region 

181 res = apdb.getDiaObjects(region) 

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

183 

184 # get sources by region 

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

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

187 

188 # get sources by object ID, empty object list 

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

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

191 

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

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

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

195 

196 def test_storeObjects(self) -> None: 

197 """Store and retrieve DiaObjects.""" 

198 

199 # don't care about sources. 

200 config = self.make_config() 

201 apdb = make_apdb(config) 

202 apdb.makeSchema() 

203 

204 region = self.make_region() 

205 visit_time = self.visit_time 

206 

207 # make catalog with Objects 

208 catalog = makeObjectCatalog(region, 100, visit_time, **self.extra_object_columns) 

209 

210 # store catalog 

211 apdb.store(visit_time, catalog) 

212 

213 # read it back and check sizes 

214 res = apdb.getDiaObjects(region) 

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

216 

217 def test_objectHistory(self) -> None: 

218 """Store and retrieve DiaObject history.""" 

219 

220 # don't care about sources. 

221 config = self.make_config() 

222 apdb = make_apdb(config) 

223 apdb.makeSchema() 

224 

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

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

227 visit_time = [ 

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

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

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

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

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

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

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

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

236 ] 

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

238 

239 nobj = 100 

240 catalog1 = makeObjectCatalog(region1, nobj, visit_time[0], **self.extra_object_columns) 

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

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

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

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

245 catalog2 = makeObjectCatalog( 

246 region2, nobj, visit_time[1], start_id=nobj*2, **self.extra_object_columns 

247 ) 

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

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

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

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

252 

253 # read it back and check sizes 

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

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

256 

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

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

259 

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

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

262 

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

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

265 

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

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

268 

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

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

271 

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

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

274 

275 res = apdb.getDiaObjectsHistory( 

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

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

278 ) 

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

280 

281 res = apdb.getDiaObjectsHistory( 

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

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

284 ) 

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

286 

287 res = apdb.getDiaObjectsHistory( 

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

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

290 ) 

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

292 

293 res = apdb.getDiaObjectsHistory( 

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

295 ) 

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

297 

298 res = apdb.getDiaObjectsHistory( 

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

300 ) 

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

302 

303 res = apdb.getDiaObjectsHistory( 

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

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

306 region1, 

307 ) 

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

309 

310 def test_storeSources(self) -> None: 

311 """Store and retrieve DiaSources.""" 

312 config = self.make_config() 

313 apdb = make_apdb(config) 

314 apdb.makeSchema() 

315 

316 region = self.make_region() 

317 visit_time = self.visit_time 

318 

319 # have to store Objects first 

320 objects = makeObjectCatalog(region, 100, visit_time, **self.extra_object_columns) 

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

322 sources = makeSourceCatalog(objects, visit_time) 

323 

324 # save the objects and sources 

325 apdb.store(visit_time, objects, sources) 

326 

327 # read it back, no ID filtering 

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

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

330 

331 # read it back and filter by ID 

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

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

334 

335 # read it back to get schema 

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

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

338 

339 def test_sourceHistory(self) -> None: 

340 """Store and retrieve DiaSource history.""" 

341 

342 # don't care about sources. 

343 config = self.make_config() 

344 apdb = make_apdb(config) 

345 apdb.makeSchema() 

346 visit_time = self.visit_time 

347 

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

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

350 nobj = 100 

351 objects1 = makeObjectCatalog(region1, nobj, visit_time, **self.extra_object_columns) 

352 objects2 = makeObjectCatalog(region2, nobj, visit_time, start_id=nobj*2, **self.extra_object_columns) 

353 

354 visits = [ 

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

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

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

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

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

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

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

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

363 ] 

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

365 

366 start_id = 0 

367 for visit_time, objects in visits: 

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

369 apdb.store(visit_time, objects, sources) 

370 start_id += nobj 

371 

372 # read it back and check sizes 

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

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

375 

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

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

378 

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

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

381 

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

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

384 

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

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

387 

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

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

390 

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

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

393 

394 res = apdb.getDiaSourcesHistory( 

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

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

397 ) 

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

399 

400 res = apdb.getDiaSourcesHistory( 

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

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

403 ) 

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

405 

406 res = apdb.getDiaSourcesHistory( 

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

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

409 ) 

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

411 

412 res = apdb.getDiaSourcesHistory( 

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

414 ) 

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

416 

417 res = apdb.getDiaSourcesHistory( 

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

419 ) 

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

421 

422 res = apdb.getDiaSourcesHistory( 

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

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

425 region1, 

426 ) 

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

428 

429 def test_storeForcedSources(self) -> None: 

430 """Store and retrieve DiaForcedSources.""" 

431 

432 config = self.make_config() 

433 apdb = make_apdb(config) 

434 apdb.makeSchema() 

435 

436 region = self.make_region() 

437 visit_time = self.visit_time 

438 

439 # have to store Objects first 

440 objects = makeObjectCatalog(region, 100, visit_time, **self.extra_object_columns) 

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

442 catalog = makeForcedSourceCatalog(objects, visit_time) 

443 

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

445 

446 # read it back and check sizes 

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

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

449 

450 # read it back to get schema 

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

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

453 

454 def test_forcedSourceHistory(self) -> None: 

455 """Store and retrieve DiaForcedSource history.""" 

456 

457 # don't care about sources. 

458 config = self.make_config() 

459 apdb = make_apdb(config) 

460 apdb.makeSchema() 

461 visit_time = self.visit_time 

462 

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

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

465 nobj = 100 

466 objects1 = makeObjectCatalog(region1, nobj, visit_time, **self.extra_object_columns) 

467 objects2 = makeObjectCatalog(region2, nobj, visit_time, start_id=nobj*2, **self.extra_object_columns) 

468 

469 visits = [ 

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

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

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

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

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

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

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

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

478 ] 

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

480 

481 start_id = 0 

482 for visit_time, objects in visits: 

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

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

485 start_id += 1 

486 

487 # read it back and check sizes 

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

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

490 

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

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

493 

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

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

496 

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

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

499 

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

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

502 

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

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

505 

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

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

508 

509 res = apdb.getDiaForcedSourcesHistory( 

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

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

512 ) 

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

514 

515 res = apdb.getDiaForcedSourcesHistory( 

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

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

518 ) 

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

520 

521 res = apdb.getDiaForcedSourcesHistory( 

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

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

524 ) 

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

526 

527 res = apdb.getDiaForcedSourcesHistory( 

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

529 ) 

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

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

532 

533 res = apdb.getDiaForcedSourcesHistory( 

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

535 ) 

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

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

538 

539 res = apdb.getDiaForcedSourcesHistory( 

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

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

542 region1, 

543 ) 

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

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

546 

547 def test_storeSSObjects(self) -> None: 

548 """Store and retrieve SSObjects.""" 

549 

550 # don't care about sources. 

551 config = self.make_config() 

552 apdb = make_apdb(config) 

553 apdb.makeSchema() 

554 

555 # make catalog with SSObjects 

556 catalog = makeSSObjectCatalog(100, flags=1) 

557 

558 # store catalog 

559 apdb.storeSSObjects(catalog) 

560 

561 # read it back and check sizes 

562 res = apdb.getSSObjects() 

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

564 

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

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

567 apdb.storeSSObjects(catalog) 

568 res = apdb.getSSObjects() 

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

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

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

572 

573 def test_reassignObjects(self) -> None: 

574 """Reassign DiaObjects.""" 

575 

576 # don't care about sources. 

577 config = self.make_config() 

578 apdb = make_apdb(config) 

579 apdb.makeSchema() 

580 

581 region = self.make_region() 

582 visit_time = self.visit_time 

583 objects = makeObjectCatalog(region, 100, visit_time, **self.extra_object_columns) 

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

585 sources = makeSourceCatalog(objects, visit_time) 

586 apdb.store(visit_time, objects, sources) 

587 

588 catalog = makeSSObjectCatalog(100) 

589 apdb.storeSSObjects(catalog) 

590 

591 # read it back and filter by ID 

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

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

594 

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

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

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

598 

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

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

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

602 

603 def test_midPointTai_src(self) -> None: 

604 """Test for time filtering of DiaSources. 

605 """ 

606 config = self.make_config() 

607 apdb = make_apdb(config) 

608 apdb.makeSchema() 

609 

610 region = self.make_region() 

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

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

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

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

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

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

617 

618 objects = makeObjectCatalog(region, 100, visit_time0, **self.extra_object_columns) 

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

620 sources = makeSourceCatalog(objects, src_time1, 0) 

621 apdb.store(src_time1, objects, sources) 

622 

623 sources = makeSourceCatalog(objects, src_time2, 100) 

624 apdb.store(src_time2, objects, sources) 

625 

626 # reading at time of last save should read all 

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

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

629 

630 # one second before 12 months 

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

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

633 

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

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

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

637 

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

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

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

641 

642 def test_midPointTai_fsrc(self) -> None: 

643 """Test for time filtering of DiaForcedSources. 

644 """ 

645 config = self.make_config() 

646 apdb = make_apdb(config) 

647 apdb.makeSchema() 

648 

649 region = self.make_region() 

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

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

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

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

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

655 

656 objects = makeObjectCatalog(region, 100, visit_time0, **self.extra_object_columns) 

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

658 sources = makeForcedSourceCatalog(objects, src_time1, 1) 

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

660 

661 sources = makeForcedSourceCatalog(objects, src_time2, 2) 

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

663 

664 # reading at time of last save should read all 

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

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

667 

668 # one second before 12 months 

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

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

671 

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

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

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

675 

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

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

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