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

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

318 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 # 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 n_ssobj_columns = 81 

57 

58 @abstractmethod 

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

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

61 raise NotImplementedError() 

62 

63 @abstractmethod 

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

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

66 raise NotImplementedError() 

67 

68 @abstractmethod 

69 def getDiaObjects_table(self) -> ApdbTables: 

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

71 raise NotImplementedError() 

72 

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

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

75 pointing_v = UnitVector3d(*xyz) 

76 fov = 0.05 # radians 

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

78 return region 

79 

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

81 """Validate catalog type and size 

82 

83 Parameters 

84 ---------- 

85 catalog : `object` 

86 Expected type of this is ``type``. 

87 rows : int 

88 Expected number of rows in a catalog. 

89 table : `ApdbTables` 

90 APDB table type. 

91 """ 

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

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

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

95 

96 def test_makeSchema(self) -> None: 

97 """Test for makeing APDB schema.""" 

98 config = self.make_config() 

99 apdb = make_apdb(config) 

100 

101 apdb.makeSchema() 

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

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

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

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

106 

107 def test_empty_gets(self) -> None: 

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

109 

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

111 checking that code is not broken. 

112 """ 

113 

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

115 config = self.make_config() 

116 apdb = make_apdb(config) 

117 apdb.makeSchema() 

118 

119 region = self.make_region() 

120 visit_time = self.visit_time 

121 

122 res: Optional[pandas.DataFrame] 

123 

124 # get objects by region 

125 res = apdb.getDiaObjects(region) 

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

127 

128 # get sources by region 

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

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

131 

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

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

134 

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

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

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

138 

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

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

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

142 

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

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

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

146 

147 # get sources by region 

148 if self.fsrc_requires_id_list: 

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

150 apdb.getDiaForcedSources(region, None, visit_time) 

151 else: 

152 apdb.getDiaForcedSources(region, None, visit_time) 

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

154 

155 def test_empty_gets_0months(self) -> None: 

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

157 

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

159 """ 

160 

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

162 config = self.make_config(read_sources_months=0, 

163 read_forced_sources_months=0) 

164 apdb = make_apdb(config) 

165 apdb.makeSchema() 

166 

167 region = self.make_region() 

168 visit_time = self.visit_time 

169 

170 res: Optional[pandas.DataFrame] 

171 

172 # get objects by region 

173 res = apdb.getDiaObjects(region) 

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

175 

176 # get sources by region 

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

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

179 

180 # get sources by object ID, empty object list 

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

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

183 

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

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

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

187 

188 def test_storeObjects(self) -> None: 

189 """Store and retrieve DiaObjects.""" 

190 

191 # don't care about sources. 

192 config = self.make_config() 

193 apdb = make_apdb(config) 

194 apdb.makeSchema() 

195 

196 region = self.make_region() 

197 visit_time = self.visit_time 

198 

199 # make catalog with Objects 

200 catalog = makeObjectCatalog(region, 100) 

201 

202 # store catalog 

203 apdb.store(visit_time, catalog) 

204 

205 # read it back and check sizes 

206 res = apdb.getDiaObjects(region) 

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

208 

209 def test_objectHistory(self) -> None: 

210 """Store and retrieve DiaObject history.""" 

211 

212 # don't care about sources. 

213 config = self.make_config() 

214 apdb = make_apdb(config) 

215 apdb.makeSchema() 

216 

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

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

219 visit_time = [ 

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

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

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

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

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

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

226 ] 

227 

228 nobj = 100 

229 catalog1 = makeObjectCatalog(region1, nobj) 

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

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

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

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

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

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

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

237 

238 # read it back and check sizes 

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

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

241 

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

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

244 

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

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

247 

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

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

250 

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

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

253 

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

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

256 

257 res = apdb.getDiaObjectsHistory( 

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

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

260 ) 

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

262 

263 res = apdb.getDiaObjectsHistory( 

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

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

266 ) 

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

268 

269 res = apdb.getDiaObjectsHistory( 

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

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

272 ) 

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

274 

275 res = apdb.getDiaObjectsHistory(DateTime("2021-01-01T00:00:00", DateTime.TAI), region=region1) 

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

277 

278 res = apdb.getDiaObjectsHistory(DateTime("2021-01-01T00:03:00", DateTime.TAI), region=region2) 

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

280 

281 res = apdb.getDiaObjectsHistory( 

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

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

284 region1, 

285 ) 

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

287 

288 def test_storeSources(self) -> None: 

289 """Store and retrieve DiaSources.""" 

290 config = self.make_config() 

291 apdb = make_apdb(config) 

292 apdb.makeSchema() 

293 

294 region = self.make_region() 

295 visit_time = self.visit_time 

296 

297 # have to store Objects first 

298 objects = makeObjectCatalog(region, 100) 

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

300 sources = makeSourceCatalog(objects, visit_time) 

301 

302 # save the objects and sources 

303 apdb.store(visit_time, objects, sources) 

304 

305 # read it back, no ID filtering 

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

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

308 

309 # read it back and filter by ID 

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

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

312 

313 # read it back to get schema 

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

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

316 

317 def test_sourceHistory(self) -> None: 

318 """Store and retrieve DiaSource history.""" 

319 

320 # don't care about sources. 

321 config = self.make_config() 

322 apdb = make_apdb(config) 

323 apdb.makeSchema() 

324 

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

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

327 nobj = 100 

328 objects1 = makeObjectCatalog(region1, nobj) 

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

330 

331 visits = [ 

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

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

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

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

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

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

338 ] 

339 

340 start_id = 0 

341 for visit_time, objects in visits: 

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

343 apdb.store(visit_time, objects, sources) 

344 start_id += nobj 

345 

346 # read it back and check sizes 

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

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

349 

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

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

352 

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

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

355 

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

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

358 

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

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

361 

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

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

364 

365 res = apdb.getDiaSourcesHistory( 

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

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

368 ) 

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

370 

371 res = apdb.getDiaSourcesHistory( 

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

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

374 ) 

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

376 

377 res = apdb.getDiaSourcesHistory( 

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

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

380 ) 

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

382 

383 res = apdb.getDiaSourcesHistory(DateTime("2021-01-01T00:00:00", DateTime.TAI), region=region1) 

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

385 

386 res = apdb.getDiaSourcesHistory(DateTime("2021-01-01T00:03:00", DateTime.TAI), region=region2) 

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

388 

389 res = apdb.getDiaSourcesHistory( 

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

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

392 region1, 

393 ) 

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

395 

396 def test_storeForcedSources(self) -> None: 

397 """Store and retrieve DiaForcedSources.""" 

398 

399 config = self.make_config() 

400 apdb = make_apdb(config) 

401 apdb.makeSchema() 

402 

403 region = self.make_region() 

404 visit_time = self.visit_time 

405 

406 # have to store Objects first 

407 objects = makeObjectCatalog(region, 100) 

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

409 catalog = makeForcedSourceCatalog(objects, visit_time) 

410 

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

412 

413 # read it back and check sizes 

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

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

416 

417 # read it back to get schema 

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

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

420 

421 def test_forcedSourceHistory(self) -> None: 

422 """Store and retrieve DiaForcedSource history.""" 

423 

424 # don't care about sources. 

425 config = self.make_config() 

426 apdb = make_apdb(config) 

427 apdb.makeSchema() 

428 

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

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

431 nobj = 100 

432 objects1 = makeObjectCatalog(region1, nobj) 

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

434 

435 visits = [ 

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

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

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

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

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

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

442 ] 

443 

444 start_id = 0 

445 for visit_time, objects in visits: 

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

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

448 start_id += 1 

449 

450 # read it back and check sizes 

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

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

453 

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

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

456 

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

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

459 

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

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

462 

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

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

465 

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

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

468 

469 res = apdb.getDiaForcedSourcesHistory( 

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

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

472 ) 

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

474 

475 res = apdb.getDiaForcedSourcesHistory( 

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

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

478 ) 

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

480 

481 res = apdb.getDiaForcedSourcesHistory( 

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

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

484 ) 

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

486 

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

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

489 

490 res = apdb.getDiaForcedSourcesHistory(DateTime("2021-01-01T00:03:00", DateTime.TAI), region=region2) 

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

492 

493 res = apdb.getDiaForcedSourcesHistory( 

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

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

496 region1, 

497 ) 

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

499 

500 def test_storeSSObjects(self) -> None: 

501 """Store and retrieve SSObjects.""" 

502 

503 # don't care about sources. 

504 config = self.make_config() 

505 apdb = make_apdb(config) 

506 apdb.makeSchema() 

507 

508 # make catalog with SSObjects 

509 catalog = makeSSObjectCatalog(100, flags=1) 

510 

511 # store catalog 

512 apdb.storeSSObjects(catalog) 

513 

514 # read it back and check sizes 

515 res = apdb.getSSObjects() 

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

517 

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

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

520 apdb.storeSSObjects(catalog) 

521 res = apdb.getSSObjects() 

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

523 self.assertEqual(len(res[res["flags"] == 1]), 50) 

524 self.assertEqual(len(res[res["flags"] == 2]), 100) 

525 

526 def test_reassignObjects(self) -> None: 

527 """Reassign DiaObjects.""" 

528 

529 # don't care about sources. 

530 config = self.make_config() 

531 apdb = make_apdb(config) 

532 apdb.makeSchema() 

533 

534 region = self.make_region() 

535 visit_time = self.visit_time 

536 objects = makeObjectCatalog(region, 100) 

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

538 sources = makeSourceCatalog(objects, visit_time) 

539 apdb.store(visit_time, objects, sources) 

540 

541 catalog = makeSSObjectCatalog(100) 

542 apdb.storeSSObjects(catalog) 

543 

544 # read it back and filter by ID 

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

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

547 

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

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

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

551 

552 def test_midPointTai_src(self) -> None: 

553 """Test for time filtering of DiaSources. 

554 """ 

555 config = self.make_config() 

556 apdb = make_apdb(config) 

557 apdb.makeSchema() 

558 

559 region = self.make_region() 

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

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

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

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

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

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

566 

567 objects = makeObjectCatalog(region, 100) 

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

569 sources = makeSourceCatalog(objects, src_time1, 0) 

570 apdb.store(src_time1, objects, sources) 

571 

572 sources = makeSourceCatalog(objects, src_time2, 100) 

573 apdb.store(src_time2, objects, sources) 

574 

575 # reading at time of last save should read all 

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

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

578 

579 # one second before 12 months 

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

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

582 

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

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

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

586 

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

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

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

590 

591 def test_midPointTai_fsrc(self) -> None: 

592 """Test for time filtering of DiaForcedSources. 

593 """ 

594 config = self.make_config() 

595 apdb = make_apdb(config) 

596 apdb.makeSchema() 

597 

598 region = self.make_region() 

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

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

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

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

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

604 

605 objects = makeObjectCatalog(region, 100) 

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

607 sources = makeForcedSourceCatalog(objects, src_time1, 1) 

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

609 

610 sources = makeForcedSourceCatalog(objects, src_time2, 2) 

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

612 

613 # reading at time of last save should read all 

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

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

616 

617 # one second before 12 months 

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

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

620 

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

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

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

624 

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

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

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