Coverage for python/lsst/dax/apdb/tests/

1# This file is part of dax_apdb.

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.

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/>.

22from __future__ import annotations

24__all__ = ["ApdbSchemaUpdateTest", "ApdbTest", "update_schema_yaml"]

26import contextlib

27import logging.config

28import os

29import tempfile

30from abc import ABC, abstractmethod

31from collections.abc import Iterator

32from tempfile import TemporaryDirectory

33from typing import TYPE_CHECKING, Any

35import astropy.time

36import felis.datamodel

37import pandas

38import yaml

40from lsst.sphgeom import Angle, Circle, LonLat, Region, UnitVector3d

42from .. import (

43 Apdb,

44 ApdbConfig,

45 ApdbReassignDiaSourceToSSObjectRecord,

46 ApdbReplica,

47 ApdbTableData,

48 ApdbTables,

49 ApdbUpdateRecord,

50 ApdbWithdrawDiaSourceRecord,

51 DiaObjectId,

52 DiaSourceId,

53 IncompatibleVersionError,

54 ReplicaChunk,

55 VersionTuple,

56)

57from .data_factory import (

58 makeForcedSourceCatalog,

59 makeObjectCatalog,

60 makeSourceCatalog,

61 makeTimestamp,

62 makeTimestampColumn,

63)

64from .utils import TestCaseMixin

66if TYPE_CHECKING:

67 from ..pixelization import Pixelization

70# Optionally configure logging from a config file.

71if log_config := os.environ.get("DAX_APDB_TEST_LOG_CONFIG"): 71 ↛ 72line 71 didn't jump to line 72 because the condition on line 71 was never true

72 logging.config.fileConfig(log_config)

75def _make_region(xyz: tuple[float, float, float] = (1.0, 1.0, -1.0)) -> Region:

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

77 pointing_v = UnitVector3d(*xyz)

78 fov = 0.0013 # radians

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

80 return region

83@contextlib.contextmanager

84def update_schema_yaml(

85 schema_file: str,

86 drop_metadata: bool = False,

87 version: str | None = None,

88) -> Iterator[str]:

89 """Update schema definition and return name of the new schema file.

91 Parameters

92 ----------

93 schema_file : `str`

94 Path for the existing YAML file with APDB schema.

95 drop_metadata : `bool`

96 If `True` then remove metadata table from the list of tables.

97 version : `str` or `None`

98 If non-empty string then set schema version to this string, if empty

99 string then remove schema version from config, if `None` - don't change

100 the version in config.

101

102 Yields

103 ------

104 Path for the updated configuration file.

105 """

106 with open(schema_file) as yaml_stream:

107 schemas_list = list(yaml.load_all(yaml_stream, Loader=yaml.SafeLoader))

108 # Edit YAML contents.

109 for schema in schemas_list:

110 # Optionally drop metadata table.

111 if drop_metadata:

112 schema["tables"] = [table for table in schema["tables"] if table["name"] != "metadata"]

113 if version is not None:

114 if version == "":

115 del schema["version"]

116 else:

117 schema["version"] = version

118

119 with TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:

120 output_path = os.path.join(tmpdir, "schema.yaml")

121 with open(output_path, "w") as yaml_stream:

122 yaml.dump_all(schemas_list, stream=yaml_stream)

123 yield output_path

124

125

126class ApdbTest(TestCaseMixin, ABC):

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

128 implementation.

129

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

131 calls various assert methods.

132 """

133

134 visit_time = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")

135

136 processing_time = astropy.time.Time("2021-01-01T12:00:00", format="isot", scale="tai")

137

138 fsrc_requires_id_list = False

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

140

141 enable_replica: bool = False

142 """Set to true when support for replication is configured"""

143

144 use_mjd: bool = True

145 """If True then timestamp columns are MJD TAI."""

146

147 extra_chunk_columns = 1

148 """Number of additional columns in chunk tables."""

149

150 meta_row_count = 3

151 """Initial row count in metadata table."""

152

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

154 table_column_count = {

155 ApdbTables.DiaObject: 8,

156 ApdbTables.DiaObjectLast: 6,

157 ApdbTables.DiaSource: 12,

158 ApdbTables.DiaForcedSource: 8,

159 ApdbTables.SSObject: 3,

160 }

161

162 @abstractmethod

163 def make_instance(self, **kwargs: Any) -> ApdbConfig:

164 """Make database instance and return configuration for it."""

165 raise NotImplementedError()

166

167 @abstractmethod

168 def getDiaObjects_table(self) -> ApdbTables:

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

170 raise NotImplementedError()

171

172 @abstractmethod

173 def pixelization(self, config: ApdbConfig) -> Pixelization:

174 """Return pixelization used by implementation."""

175 raise NotImplementedError()

176

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

178 """Validate catalog type and size

179

180 Parameters

181 ----------

182 catalog : `object`

183 Expected type of this is ``pandas.DataFrame``.

184 rows : `int`

185 Expected number of rows in a catalog.

186 table : `ApdbTables`

187 APDB table type.

188 """

189 self.assertIsInstance(catalog, pandas.DataFrame)

190 self.assertEqual(catalog.shape[0], rows)

191 self.assertEqual(catalog.shape[1], self.table_column_count[table])

192

193 def assert_table_data(self, catalog: Any, rows: int, table: ApdbTables) -> None:

194 """Validate catalog type and size

195

196 Parameters

197 ----------

198 catalog : `object`

199 Expected type of this is `ApdbTableData`.

200 rows : `int`

201 Expected number of rows in a catalog.

202 table : `ApdbTables`

203 APDB table type.

204 extra_columns : `int`

205 Count of additional columns expected in ``catalog``.

206 """

207 self.assertIsInstance(catalog, ApdbTableData)

208 n_rows = sum(1 for row in catalog.rows())

209 self.assertEqual(n_rows, rows)

210 # One extra column for replica chunk id

211 self.assertEqual(

212 len(catalog.column_names()), self.table_column_count[table] + self.extra_chunk_columns

213 )

214

215 def assert_column_types(self, catalog: Any, types: dict[str, felis.datamodel.DataType]) -> None:

216 column_defs = dict(catalog.column_defs())

217 for column, datatype in types.items():

218 self.assertEqual(column_defs[column], datatype)

219

220 def make_region(self, xyz: tuple[float, float, float] = (1.0, 1.0, -1.0)) -> Region:

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

222 return _make_region(xyz)

223

224 def test_makeSchema(self) -> None:

225 """Test for making APDB schema."""

226 config = self.make_instance()

227 apdb = Apdb.from_config(config)

228

229 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject))

230 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObjectLast))

231 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaSource))

232 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaForcedSource))

233 self.assertIsNotNone(apdb.tableDef(ApdbTables.metadata))

234 self.assertIsNotNone(apdb.tableDef(ApdbTables.SSObject))

235 self.assertIsNotNone(apdb.tableDef(ApdbTables.SSSource))

236 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject_To_Object_Match))

237

238 # Test from_uri factory method with the same config.

239 with tempfile.NamedTemporaryFile() as tmpfile:

240 config.save(tmpfile.name)

241 apdb = Apdb.from_uri(tmpfile.name)

242

243 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject))

244 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObjectLast))

245 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaSource))

246 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaForcedSource))

247 self.assertIsNotNone(apdb.tableDef(ApdbTables.metadata))

248 self.assertIsNotNone(apdb.tableDef(ApdbTables.SSObject))

249 self.assertIsNotNone(apdb.tableDef(ApdbTables.SSSource))

250 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject_To_Object_Match))

251

252 def test_empty_gets(self) -> None:

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

254

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

256 checking that code is not broken.

257 """

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

259 config = self.make_instance()

260 apdb = Apdb.from_config(config)

261

262 region = self.make_region()

263 visit_time = self.visit_time

264

265 res: pandas.DataFrame | None

266

267 # get objects by region

268 res = apdb.getDiaObjects(region)

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

270

271 # get sources by region

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

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

274

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

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

277

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

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

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

281

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

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

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

285

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

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

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

289

290 # data_factory's ccdVisitId generation corresponds to (1, 1)

291 res = apdb.containsVisitDetector(visit=1, detector=1, region=region, visit_time=visit_time)

292 self.assertFalse(res)

293

294 # get sources by region

295 if self.fsrc_requires_id_list:

296 with self.assertRaises(NotImplementedError):

297 apdb.getDiaForcedSources(region, None, visit_time)

298 else:

299 res = apdb.getDiaForcedSources(region, None, visit_time)

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

301

302 def test_empty_gets_0months(self) -> None:

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

304

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

306 """

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

308 config = self.make_instance(read_sources_months=0, read_forced_sources_months=0)

309 apdb = Apdb.from_config(config)

310

311 region = self.make_region()

312 visit_time = self.visit_time

313

314 res: pandas.DataFrame | None

315

316 # get objects by region

317 res = apdb.getDiaObjects(region)

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

319

320 # get sources by region

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

322 self.assertIs(res, None)

323

324 # get sources by object ID, empty object list

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

326 self.assertIs(res, None)

327

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

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

330 self.assertIs(res, None)

331

332 # Database is empty, no images exist.

333 res = apdb.containsVisitDetector(visit=1, detector=1, region=region, visit_time=visit_time)

334 self.assertFalse(res)

335

336 def test_storeObjects(self) -> None:

337 """Store and retrieve DiaObjects."""

338 # don't care about sources.

339 config = self.make_instance()

340 apdb = Apdb.from_config(config)

341

342 region = self.make_region()

343 visit_time = self.visit_time

344

345 # make catalog with Objects

346 catalog = makeObjectCatalog(region, 100)

347

348 # store catalog

349 apdb.store(visit_time, catalog)

350

351 # read it back and check sizes

352 res = apdb.getDiaObjects(region)

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

354

355 # TODO: test apdb.contains with generic implementation from DM-41671

356

357 def test_storeObjects_empty(self) -> None:

358 """Test calling storeObject when there are no objects: see DM-43270."""

359 config = self.make_instance()

360 apdb = Apdb.from_config(config)

361 region = self.make_region()

362 visit_time = self.visit_time

363 # make catalog with no Objects

364 catalog = makeObjectCatalog(region, 0)

365

366 with self.assertLogs("lsst.dax.apdb", level="DEBUG") as cm:

367 apdb.store(visit_time, catalog)

368 self.assertIn("No objects", "\n".join(cm.output))

369

370 def test_storeMovingObject(self) -> None:

371 """Store and retrieve DiaObject which changes its position."""

372 # don't care about sources.

373 config = self.make_instance()

374 apdb = Apdb.from_config(config)

375 pixelization = self.pixelization(config)

376

377 lon_deg, lat_deg = 0.0, 0.0

378 lonlat1 = LonLat.fromDegrees(lon_deg - 1.0, lat_deg)

379 lonlat2 = LonLat.fromDegrees(lon_deg + 1.0, lat_deg)

380 uv1 = UnitVector3d(lonlat1)

381 uv2 = UnitVector3d(lonlat2)

382

383 # Check that they fall into different pixels.

384 self.assertNotEqual(pixelization.pixel(uv1), pixelization.pixel(uv2))

385

386 # Store one object at two different positions.

387 visit_time1 = self.visit_time

388 catalog1 = makeObjectCatalog(lonlat1, 1)

389 apdb.store(visit_time1, catalog1)

390

391 visit_time2 = visit_time1 + astropy.time.TimeDelta(120.0, format="sec")

392 catalog1 = makeObjectCatalog(lonlat2, 1)

393 apdb.store(visit_time2, catalog1)

394

395 # Make region covering both points.

396 region = Circle(UnitVector3d(LonLat.fromDegrees(lon_deg, lat_deg)), Angle.fromDegrees(1.1))

397 self.assertTrue(region.contains(uv1))

398 self.assertTrue(region.contains(uv2))

399

400 # Read it back, must return the latest one.

401 res = apdb.getDiaObjects(region)

402 self.assert_catalog(res, 1, self.getDiaObjects_table())

403

404 def test_storeSources(self) -> None:

405 """Store and retrieve DiaSources."""

406 config = self.make_instance()

407 apdb = Apdb.from_config(config)

408

409 region = self.make_region()

410 visit_time = self.visit_time

411

412 # have to store Objects first

413 objects = makeObjectCatalog(region, 100)

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

415 sources = makeSourceCatalog(objects, visit_time, use_mjd=self.use_mjd)

416

417 # save the objects and sources

418 apdb.store(visit_time, objects, sources)

419

420 # read it back, no ID filtering

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

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

423

424 # read it back and filter by ID

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

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

427

428 # read it back to get schema

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

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

431

432 # test if a visit is present

433 # data_factory's ccdVisitId generation corresponds to (1, 1)

434 res = apdb.containsVisitDetector(visit=1, detector=1, region=region, visit_time=visit_time)

435 self.assertTrue(res)

436 # non-existent image

437 res = apdb.containsVisitDetector(visit=2, detector=42, region=region, visit_time=visit_time)

438 self.assertFalse(res)

439

440 def test_storeForcedSources(self) -> None:

441 """Store and retrieve DiaForcedSources."""

442 config = self.make_instance()

443 apdb = Apdb.from_config(config)

444

445 region = self.make_region()

446 visit_time = self.visit_time

447

448 # have to store Objects first

449 objects = makeObjectCatalog(region, 100)

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

451 catalog = makeForcedSourceCatalog(objects, visit_time, use_mjd=self.use_mjd)

452

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

454

455 # read it back and check sizes

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

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

458

459 # read it back to get schema

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

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

462

463 # data_factory's ccdVisitId generation corresponds to (1, 1)

464 res = apdb.containsVisitDetector(visit=1, detector=1, region=region, visit_time=visit_time)

465 self.assertTrue(res)

466 # non-existent image

467 res = apdb.containsVisitDetector(visit=2, detector=42, region=region, visit_time=visit_time)

468 self.assertFalse(res)

469

470 def test_timestamps(self) -> None:

471 """Check that timestamp return type is as expected."""

472 config = self.make_instance()

473 apdb = Apdb.from_config(config)

474

475 region = self.make_region()

476 visit_time = self.visit_time

477

478 # Cassandra has a millisecond precision, so subtract 1ms to allow for

479 # truncated returned values.

480 time_before = makeTimestamp(self.processing_time, self.use_mjd, -1)

481 objects = makeObjectCatalog(region, 100)

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

483 catalog = makeForcedSourceCatalog(

484 objects, visit_time, processing_time=self.processing_time, use_mjd=self.use_mjd

485 )

486 time_after = makeTimestamp(self.processing_time, self.use_mjd)

487

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

489

490 # read it back and check sizes

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

492 assert res is not None

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

494

495 time_processed_column = makeTimestampColumn("time_processed", self.use_mjd)

496 self.assertIn(time_processed_column, res.dtypes)

497 dtype = res.dtypes[time_processed_column]

498 timestamp_type_names = ("float64",) if self.use_mjd else ("datetime64[us]", "datetime64[ns]")

499 self.assertIn(dtype.name, timestamp_type_names)

500 # Verify that returned time is sensible.

501 self.assertTrue(all(time_before <= dt <= time_after for dt in res[time_processed_column]))

502

503 def test_getDiaObjectsForDedup(self) -> None:

504 """Test getDiaObjectsForDedup() method."""

505 config = self.make_instance()

506 apdb = Apdb.from_config(config)

507

508 region1 = self.make_region((1.0, 1.0, -1.0))

509 region2 = self.make_region((-1.0, 1.0, -1.0))

510 region3 = self.make_region((-1.0, -1.0, -1.0))

511 nobj = 100

512 objects1 = makeObjectCatalog(region1, nobj)

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

514 objects3 = makeObjectCatalog(region3, nobj, start_id=nobj * 4)

515

516 visits = [

517 (astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai"), objects1),

518 (astropy.time.Time("2021-01-01T00:10:00", format="isot", scale="tai"), objects2),

519 (astropy.time.Time("2021-01-01T00:20:00", format="isot", scale="tai"), objects3),

520 ]

521

522 for visit_time, objects in visits:

523 apdb.store(visit_time, objects)

524

525 catalog = apdb.getDiaObjectsForDedup()

526 self.assertEqual(len(catalog), 300)

527

528 catalog = apdb.getDiaObjectsForDedup(visits[0][0])

529 self.assertEqual(len(catalog), 300)

530

531 catalog = apdb.getDiaObjectsForDedup(visits[1][0])

532 self.assertEqual(len(catalog), 200)

533

534 catalog = apdb.getDiaObjectsForDedup(visits[2][0])

535 self.assertEqual(len(catalog), 100)

536

537 time = astropy.time.Time("2021-01-01T00:30:00", format="isot", scale="tai")

538 catalog = apdb.getDiaObjectsForDedup(time)

539 self.assertEqual(len(catalog), 0)

540

541 def test_getDiaSourcesForDiaObjects(self) -> None:

542 """Test getDiaSourcesForDiaObjects() method."""

543 config = self.make_instance()

544 apdb = Apdb.from_config(config)

545 # Monkey-patch APDB instance to set current time.

546 apdb._current_time = lambda: self.processing_time # type: ignore[method-assign]

547

548 region1 = self.make_region((1.0, 1.0, -1.0))

549 region2 = self.make_region((-1.0, 1.0, -1.0))

550 region3 = self.make_region((-1.0, -1.0, -1.0))

551 nobj = 100

552 objects1 = makeObjectCatalog(region1, nobj)

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

554 objects3 = makeObjectCatalog(region3, nobj, start_id=nobj * 4)

555

556 visits = [

557 (astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai"), objects1),

558 (astropy.time.Time("2021-01-01T00:10:00", format="isot", scale="tai"), objects2),

559 (astropy.time.Time("2021-01-01T00:20:00", format="isot", scale="tai"), objects3),

560 ]

561

562 start_id = 1_000_000

563 for visit_time, objects in visits:

564 sources = makeSourceCatalog(objects, visit_time, start_id=start_id, use_mjd=self.use_mjd)

565 apdb.store(visit_time, objects, sources)

566 start_id += 1_000_000

567

568 # Take a small number of objects from different regions.

569 object_ids = [

570 DiaObjectId.from_named_tuple(next(objects1.itertuples())),

571 DiaObjectId.from_named_tuple(next(objects2.itertuples())),

572 DiaObjectId.from_named_tuple(next(objects3.itertuples())),

573 ]

574

575 catalog = apdb.getDiaSourcesForDiaObjects(object_ids, visits[0][0])

576 self.assertEqual(len(catalog), 3)

577 self.assertEqual(set(catalog["diaObjectId"]), {1, 200, 400})

578 self.assertEqual(set(catalog["diaSourceId"]), {1_000_000, 2_000_000, 3_000_000})

579

580 catalog = apdb.getDiaSourcesForDiaObjects(object_ids, visits[2][0])

581 self.assertEqual(len(catalog), 1)

582 self.assertEqual(set(catalog["diaObjectId"]), {400})

583 self.assertEqual(set(catalog["diaSourceId"]), {3_000_000})

584

585 def test_reassignDiaSourcesToDiaObjects(self) -> None:

586 """Test reassignDiaSourcesToDiaObjects() method."""

587 config = self.make_instance()

588 apdb = Apdb.from_config(config)

589 apdb._current_time = lambda: self.processing_time # type: ignore[method-assign]

590 apdb_replica = ApdbReplica.from_config(config)

591

592 visit_time = self.visit_time

593 lonlat1 = LonLat.fromDegrees(0.0, 0.0)

594 lonlat2 = LonLat.fromDegrees(180.0, 0.0)

595 # regons around lonlat1/2

596 region1 = self.make_region(xyz=(1.0, 0.0, 0.0))

597 region2 = self.make_region(xyz=(-1.0, 0.0, 0.0))

598

599 # Store 3 objects and sources at the same position in each region.

600 objects = makeObjectCatalog(lonlat1, 3, start_id=100)

601 sources = makeSourceCatalog(objects, visit_time, start_id=1000, use_mjd=self.use_mjd)

602 apdb.store(visit_time, objects, sources)

603

604 objects = makeObjectCatalog(lonlat2, 3, start_id=200)

605 sources = makeSourceCatalog(objects, visit_time, start_id=2000, use_mjd=self.use_mjd)

606 apdb.store(visit_time, objects, sources)

607

608 # check that everything as we think it is.

609 objects = apdb.getDiaObjects(region1)

610 self.assertEqual(set(objects["diaObjectId"]), {100, 101, 102})

611 self.assertEqual(list(objects["nDiaSources"]), [1, 1, 1])

612 sources = apdb.getDiaSources(region1, [100, 101, 102], visit_time)

613 assert sources is not None

614 self.assertEqual(set(sources["diaSourceId"]), {1000, 1001, 1002})

615 self.assertEqual(set(sources["diaObjectId"]), {100, 101, 102})

616

617 # Reassign sources in region1 and increment/decrement nDiaSources.

618 midpointMjdTai = visit_time.tai.mjd

619 reassign = {

620 DiaSourceId(diaSourceId=1001, ra=0.0, dec=0.0, midpointMjdTai=midpointMjdTai): 100,

621 DiaSourceId(diaSourceId=1002, ra=0.0, dec=0.0, midpointMjdTai=midpointMjdTai): 100,

622 }

623 apdb.reassignDiaSourcesToDiaObjects(reassign)

624

625 objects = apdb.getDiaObjects(region1)

626 self.assertEqual(set(objects["nDiaSources"]), {0, 3})

627 sources = apdb.getDiaSources(region1, [100], visit_time)

628 assert sources is not None

629 self.assertEqual(set(sources["diaSourceId"]), {1000, 1001, 1002})

630 self.assertEqual(set(sources["diaObjectId"]), {100})

631

632 # Reassign but do not increment/decrement nDiaSources.

633 reassign = {

634 DiaSourceId(diaSourceId=2001, ra=180.0, dec=0.0, midpointMjdTai=midpointMjdTai): 200,

635 DiaSourceId(diaSourceId=2002, ra=180.0, dec=0.0, midpointMjdTai=midpointMjdTai): 200,

636 }

637 apdb.reassignDiaSourcesToDiaObjects(

638 reassign, increment_nDiaSources=False, decrement_nDiaSources=False

639 )

640

641 objects = apdb.getDiaObjects(region2)

642 self.assertEqual(set(objects["nDiaSources"]), {1})

643 sources = apdb.getDiaSources(region2, [200], visit_time)

644 assert sources is not None

645 self.assertEqual(set(sources["diaSourceId"]), {2000, 2001, 2002})

646 self.assertEqual(set(sources["diaObjectId"]), {200})

647

648 replica_chunks = apdb_replica.getReplicaChunks()

649 if not self.enable_replica:

650 self.assertIsNone(replica_chunks)

651 else:

652 assert replica_chunks is not None

653

654 # There could be one or two chunks.

655 self.assertTrue(1 <= len(replica_chunks) <= 2)

656

657 update_records = apdb_replica.getUpdateRecordChunks([chunk.id for chunk in replica_chunks])

658 # Two reassignments for region1, three increments/decrements for

659 # that region, plus two reassignments for region2 without

660 # increments/decrements.

661 self.assertEqual(len(update_records), 2 + 3 + 2)

662

663 def test_setValidityEnd(self) -> None:

664 """Store DiaObjects and truncate validity for some."""

665 # don't care about sources.

666 config = self.make_instance()

667 apdb = Apdb.from_config(config)

668 apdb._current_time = lambda: self.processing_time # type: ignore[method-assign]

669 apdb_replica = ApdbReplica.from_config(config)

670

671 region = self.make_region()

672 visit_time = self.visit_time

673

674 # make catalog with Objects

675 catalog = makeObjectCatalog(region, 100)

676

677 # store catalog

678 apdb.store(visit_time, catalog)

679

680 # read it back and check sizes

681 res = apdb.getDiaObjects(region)

682 self.assert_catalog(res, 100, self.getDiaObjects_table())

683

684 # Select first 10 objects.

685 object_ids = [DiaObjectId.from_named_tuple(row) for row in catalog.iloc[:10].itertuples()]

686 count = apdb.setValidityEnd(object_ids, self.processing_time)

687 self.assertEqual(count, 10)

688

689 res = apdb.getDiaObjects(region)

690 self.assert_catalog(res, 90, self.getDiaObjects_table())

691

692 replica_chunks = apdb_replica.getReplicaChunks()

693 if not self.enable_replica:

694 self.assertIsNone(replica_chunks)

695 else:

696 # Check that there are 10 update records in replica tables.

697 assert replica_chunks is not None

698

699 # There could be one or two chunks.

700 self.assertTrue(1 <= len(replica_chunks) <= 2)

701

702 update_records = apdb_replica.getUpdateRecordChunks([chunk.id for chunk in replica_chunks])

703 self.assertEqual(len(update_records), 10)

704

705 # Check that empty list works.

706 count = apdb.setValidityEnd(object_ids, self.processing_time)

707 self.assertEqual(count, 0)

708

709 # Try with non-existing object.

710 object_ids = [DiaObjectId.from_named_tuple(row) for row in catalog.iloc[10:12].itertuples()]

711 object_ids += [DiaObjectId(diaObjectId=1_000_000, ra=0.0, dec=0.0)]

712 with self.assertRaises(LookupError):

713 apdb.setValidityEnd(object_ids, self.processing_time, raise_on_missing_id=True)

714

715 count = apdb.setValidityEnd(object_ids, self.processing_time)

716 self.assertEqual(count, 2)

717

718 def test_resetDedup(self) -> None:

719 """Test resetDedup method."""

720 # don't care about sources.

721 config = self.make_instance()

722 apdb = Apdb.from_config(config)

723

724 region = self.make_region()

725

726 # make catalog with Objects

727 objects = makeObjectCatalog(region, 100)

728

729 visit_time1 = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")

730 dedup_time1 = astropy.time.Time("2021-01-01T12:00:00", format="isot", scale="tai")

731 visit_time2 = astropy.time.Time("2021-01-02T00:00:00", format="isot", scale="tai")

732 dedup_time2 = astropy.time.Time("2021-01-02T12:00:00", format="isot", scale="tai")

733

734 # store catalog

735 apdb.store(visit_time1, objects)

736

737 catalog = apdb.getDiaObjectsForDedup()

738 self.assertEqual(len(catalog), 100)

739

740 catalog = apdb.getDiaObjectsForDedup(visit_time1)

741 self.assertEqual(len(catalog), 100)

742

743 apdb.resetDedup(dedup_time1)

744

745 catalog = apdb.getDiaObjectsForDedup(visit_time1)

746 self.assertEqual(len(catalog), self._count_after_reset_dedup(100))

747

748 apdb.store(visit_time2, objects)

749

750 catalog = apdb.getDiaObjectsForDedup()

751 self.assertEqual(len(catalog), 100)

752

753 catalog = apdb.getDiaObjectsForDedup(dedup_time1)

754 self.assertEqual(len(catalog), 100)

755

756 apdb.resetDedup(dedup_time2)

757

758 catalog = apdb.getDiaObjectsForDedup(dedup_time1)

759 self.assertEqual(len(catalog), self._count_after_reset_dedup(100))

760

761 catalog = apdb.getDiaObjectsForDedup()

762 self.assertEqual(len(catalog), 0)

763

764 def _count_after_reset_dedup(self, count_before: int) -> int:

765 """Return the number of rows that will be returned by

766 getDiaObjectsForDedup() after resetDedup() was called. For SQL backend

767 deduplication data comes from a regular table, and it is not removed

768 by resetDedup().

769 """

770 raise NotImplementedError()

771

772 def test_getChunks(self) -> None:

773 """Store and retrieve replica chunks."""

774 # don't care about sources.

775 config = self.make_instance()

776 apdb = Apdb.from_config(config)

777 apdb_replica = ApdbReplica.from_config(config)

778 visit_time = self.visit_time

779

780 region1 = self.make_region((1.0, 1.0, -1.0))

781 region2 = self.make_region((-1.0, -1.0, -1.0))

782 nobj = 100

783 objects1 = makeObjectCatalog(region1, nobj)

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

785

786 # With the default 10 minutes replica chunk window we should have 4

787 # records.

788 visits = [

789 (astropy.time.Time("2021-01-01T00:01:00", format="isot", scale="tai"), objects1),

790 (astropy.time.Time("2021-01-01T00:02:00", format="isot", scale="tai"), objects2),

791 (astropy.time.Time("2021-01-01T00:11:00", format="isot", scale="tai"), objects1),

792 (astropy.time.Time("2021-01-01T00:12:00", format="isot", scale="tai"), objects2),

793 (astropy.time.Time("2021-01-01T00:45:00", format="isot", scale="tai"), objects1),

794 (astropy.time.Time("2021-01-01T00:46:00", format="isot", scale="tai"), objects2),

795 (astropy.time.Time("2021-03-01T00:01:00", format="isot", scale="tai"), objects1),

796 (astropy.time.Time("2021-03-01T00:02:00", format="isot", scale="tai"), objects2),

797 ]

798

799 start_id = 0

800 for visit_time, objects in visits:

801 sources = makeSourceCatalog(objects, visit_time, start_id=start_id, use_mjd=self.use_mjd)

802 fsources = makeForcedSourceCatalog(objects, visit_time, visit=start_id, use_mjd=self.use_mjd)

803 apdb.store(visit_time, objects, sources, fsources)

804 start_id += nobj

805

806 replica_chunks = apdb_replica.getReplicaChunks()

807 if not self.enable_replica:

808 self.assertIsNone(replica_chunks)

809

810 with self.assertRaisesRegex(ValueError, "APDB is not configured for replication"):

811 apdb_replica.getTableDataChunks(ApdbTables.DiaObject, [])

812

813 else:

814 assert replica_chunks is not None

815 self.assertEqual(len(replica_chunks), 4)

816

817 with self.assertRaisesRegex(ValueError, "does not support replica chunks"):

818 apdb_replica.getTableDataChunks(ApdbTables.SSObject, [])

819

820 def _check_chunks(replica_chunks: list[ReplicaChunk], n_records: int | None = None) -> None:

821 if n_records is None:

822 n_records = len(replica_chunks) * nobj

823 res = apdb_replica.getTableDataChunks(

824 ApdbTables.DiaObject, (chunk.id for chunk in replica_chunks)

825 )

826 self.assert_table_data(res, n_records, ApdbTables.DiaObject)

827 validityStartColumn = "validityStartMjdTai" if self.use_mjd else "validityStart"

828 validityStartType = (

829 felis.datamodel.DataType.double if self.use_mjd else felis.datamodel.DataType.timestamp

830 )

831 self.assert_column_types(

832 res,

833 {

834 "apdb_replica_chunk": felis.datamodel.DataType.long,

835 "diaObjectId": felis.datamodel.DataType.long,

836 validityStartColumn: validityStartType,

837 "ra": felis.datamodel.DataType.double,

838 "dec": felis.datamodel.DataType.double,

839 "parallax": felis.datamodel.DataType.float,

840 "nDiaSources": felis.datamodel.DataType.int,

841 },

842 )

843

844 res = apdb_replica.getTableDataChunks(

845 ApdbTables.DiaSource, (chunk.id for chunk in replica_chunks)

846 )

847 self.assert_table_data(res, n_records, ApdbTables.DiaSource)

848 self.assert_column_types(

849 res,

850 {

851 "apdb_replica_chunk": felis.datamodel.DataType.long,

852 "diaSourceId": felis.datamodel.DataType.long,

853 "visit": felis.datamodel.DataType.long,

854 "detector": felis.datamodel.DataType.short,

855 },

856 )

857

858 res = apdb_replica.getTableDataChunks(

859 ApdbTables.DiaForcedSource, (chunk.id for chunk in replica_chunks)

860 )

861 self.assert_table_data(res, n_records, ApdbTables.DiaForcedSource)

862 self.assert_column_types(

863 res,

864 {

865 "apdb_replica_chunk": felis.datamodel.DataType.long,

866 "diaObjectId": felis.datamodel.DataType.long,

867 "visit": felis.datamodel.DataType.long,

868 "detector": felis.datamodel.DataType.short,

869 },

870 )

871

872 # read it back and check sizes

873 _check_chunks(replica_chunks, 800)

874 _check_chunks(replica_chunks[1:], 600)

875 _check_chunks(replica_chunks[1:-1], 400)

876 _check_chunks(replica_chunks[2:3], 200)

877 _check_chunks([])

878

879 # try to remove some of those

880 deleted_chunks = replica_chunks[:1]

881 apdb_replica.deleteReplicaChunks(chunk.id for chunk in deleted_chunks)

882

883 # All queries on deleted ids should return empty set.

884 _check_chunks(deleted_chunks, 0)

885

886 replica_chunks = apdb_replica.getReplicaChunks()

887 assert replica_chunks is not None

888 self.assertEqual(len(replica_chunks), 3)

889

890 _check_chunks(replica_chunks, 600)

891

892 def test_reassignObjects(self) -> None:

893 """Reassign DiaObjects."""

894 # don't care about sources.

895 config = self.make_instance()

896 apdb = Apdb.from_config(config)

897

898 region = self.make_region()

899 visit_time = self.visit_time

900 objects = makeObjectCatalog(region, 100)

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

902 sources = makeSourceCatalog(objects, visit_time, use_mjd=self.use_mjd)

903 apdb.store(visit_time, objects, sources)

904

905 # read it back and filter by ID

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

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

908

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

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

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

912

913 with self.assertRaisesRegex(ValueError, r"do not exist.*\D1000"):

914 apdb.reassignDiaSources(

915 {

916 1000: 1,

917 7: 3,

918 }

919 )

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

921

922 def test_storeUpdateRecord(self) -> None:

923 """Test _storeUpdateRecord() method."""

924 config = self.make_instance()

925 apdb = Apdb.from_config(config)

926

927 # Times are totally arbitrary.

928 update_time_ns1 = 2_000_000_000_000_000_000

929 update_time_ns2 = 2_000_000_001_000_000_000

930 records = [

931 ApdbReassignDiaSourceToSSObjectRecord(

932 update_time_ns=update_time_ns1,

933 update_order=0,

934 diaSourceId=1,

935 ssObjectId=1,

936 ssObjectReassocTimeMjdTai=60000.0,

937 ra=45.0,

938 dec=-45.0,

939 midpointMjdTai=60000.0,

940 ),

941 ApdbWithdrawDiaSourceRecord(

942 update_time_ns=update_time_ns1,

943 update_order=1,

944 diaSourceId=123456,

945 timeWithdrawnMjdTai=61000.0,

946 ra=45.0,

947 dec=-45.0,

948 midpointMjdTai=60000.0,

949 ),

950 ApdbReassignDiaSourceToSSObjectRecord(

951 update_time_ns=update_time_ns1,

952 update_order=3,

953 diaSourceId=2,

954 ssObjectId=3,

955 ssObjectReassocTimeMjdTai=60000.0,

956 ra=45.0,

957 dec=-45.0,

958 midpointMjdTai=60000.0,

959 ),

960 ApdbWithdrawDiaSourceRecord(

961 update_time_ns=update_time_ns2,

962 update_order=0,

963 diaSourceId=123456,

964 timeWithdrawnMjdTai=61000.0,

965 ra=45.0,

966 dec=-45.0,

967 midpointMjdTai=60000.0,

968 ),

969 ]

970

971 update_time = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")

972 chunk = ReplicaChunk.make_replica_chunk(update_time, 600)

973

974 if not self.enable_replica:

975 with self.assertRaises(TypeError):

976 self.store_update_records(apdb, records, chunk)

977 else:

978 self.store_update_records(apdb, records, chunk)

979

980 apdb_replica = ApdbReplica.from_config(config)

981 records_returned = apdb_replica.getUpdateRecordChunks([chunk.id])

982

983 # Input records are ordered, output will be ordered too.

984 self.assertEqual(records_returned, records)

985

986 @abstractmethod

987 def store_update_records(self, apdb: Apdb, records: list[ApdbUpdateRecord], chunk: ReplicaChunk) -> None:

988 """Store update records in database, must be overriden in subclass."""

989 raise NotImplementedError()

990

991 def test_midpointMjdTai_src(self) -> None:

992 """Test for time filtering of DiaSources."""

993 config = self.make_instance()

994 apdb = Apdb.from_config(config)

995

996 region = self.make_region()

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

998 src_time1 = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")

999 src_time2 = astropy.time.Time("2021-01-01T00:00:02", format="isot", scale="tai")

1000 visit_time0 = astropy.time.Time("2021-12-26T23:59:59", format="isot", scale="tai")

1001 visit_time1 = astropy.time.Time("2021-12-27T00:00:01", format="isot", scale="tai")

1002 visit_time2 = astropy.time.Time("2021-12-27T00:00:03", format="isot", scale="tai")

1003 one_sec = astropy.time.TimeDelta(1.0, format="sec")

1004

1005 objects = makeObjectCatalog(region, 100)

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

1007 sources = makeSourceCatalog(objects, src_time1, 0, use_mjd=self.use_mjd)

1008 apdb.store(src_time1, objects, sources)

1009

1010 sources = makeSourceCatalog(objects, src_time2, 100, use_mjd=self.use_mjd)

1011 apdb.store(src_time2, objects, sources)

1012

1013 # reading at time of last save should read all

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

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

1016

1017 # one second before 12 months

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

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

1020

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

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

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

1024

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

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

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

1028

1029 # Use explicit start time argument instead of 12 month window, visit

1030 # time does not matter in this case, set it to before all data.

1031 res = apdb.getDiaSources(region, oids, src_time1 - one_sec, src_time1 - one_sec)

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

1033

1034 res = apdb.getDiaSources(region, oids, src_time1 - one_sec, src_time2 - one_sec)

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

1036

1037 res = apdb.getDiaSources(region, oids, src_time1 - one_sec, src_time2 + one_sec)

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

1039

1040 def test_midpointMjdTai_fsrc(self) -> None:

1041 """Test for time filtering of DiaForcedSources."""

1042 config = self.make_instance()

1043 apdb = Apdb.from_config(config)

1044

1045 region = self.make_region()

1046 src_time1 = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")

1047 src_time2 = astropy.time.Time("2021-01-01T00:00:02", format="isot", scale="tai")

1048 visit_time0 = astropy.time.Time("2021-12-26T23:59:59", format="isot", scale="tai")

1049 visit_time1 = astropy.time.Time("2021-12-27T00:00:01", format="isot", scale="tai")

1050 visit_time2 = astropy.time.Time("2021-12-27T00:00:03", format="isot", scale="tai")

1051 one_sec = astropy.time.TimeDelta(1.0, format="sec")

1052

1053 objects = makeObjectCatalog(region, 100)

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

1055 sources = makeForcedSourceCatalog(objects, src_time1, 1, use_mjd=self.use_mjd)

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

1057

1058 sources = makeForcedSourceCatalog(objects, src_time2, 2, use_mjd=self.use_mjd)

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

1060

1061 # reading at time of last save should read all

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

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

1064

1065 # one second before 12 months

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

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

1068

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

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

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

1072

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

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

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

1076

1077 # Use explicit start time argument instead of 12 month window, visit

1078 # time does not matter in this case, set it to before all data.

1079 res = apdb.getDiaForcedSources(region, oids, src_time1 - one_sec, src_time1 - one_sec)

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

1081

1082 res = apdb.getDiaForcedSources(region, oids, src_time1 - one_sec, src_time2 - one_sec)

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

1084

1085 res = apdb.getDiaForcedSources(region, oids, src_time1 - one_sec, src_time2 + one_sec)

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

1087

1088 def test_metadata(self) -> None:

1089 """Simple test for writing/reading metadata table"""

1090 config = self.make_instance()

1091 apdb = Apdb.from_config(config)

1092 metadata = apdb.metadata

1093

1094 # APDB should write two or three metadata items with version numbers

1095 # and a frozen JSON config.

1096 self.assertFalse(metadata.empty())

1097 self.assertEqual(len(list(metadata.items())), self.meta_row_count)

1098

1099 metadata.set("meta", "data")

1100 metadata.set("data", "meta")

1101

1102 self.assertFalse(metadata.empty())

1103 self.assertTrue(set(metadata.items()) >= {("meta", "data"), ("data", "meta")})

1104

1105 with self.assertRaisesRegex(KeyError, "Metadata key 'meta' already exists"):

1106 metadata.set("meta", "data1")

1107

1108 metadata.set("meta", "data2", force=True)

1109 self.assertTrue(set(metadata.items()) >= {("meta", "data2"), ("data", "meta")})

1110

1111 self.assertTrue(metadata.delete("meta"))

1112 self.assertIsNone(metadata.get("meta"))

1113 self.assertFalse(metadata.delete("meta"))

1114

1115 self.assertEqual(metadata.get("data"), "meta")

1116 self.assertEqual(metadata.get("meta", "meta"), "meta")

1117

1118 def test_schemaVersionFromYaml(self) -> None:

1119 """Check version number handling for reading schema from YAML."""

1120 config = self.make_instance()

1121 default_schema = config.schema_file

1122 apdb = Apdb.from_config(config)

1123 self.assertEqual(apdb.schema.schemaVersion(), VersionTuple(0, 1, 1))

1124

1125 with update_schema_yaml(default_schema, version="") as schema_file:

1126 config = self.make_instance(schema_file=schema_file)

1127 apdb = Apdb.from_config(config)

1128 self.assertEqual(

1129 apdb.schema.schemaVersion(),

1130 VersionTuple(0, 1, 0),

1131 )

1132

1133 with update_schema_yaml(default_schema, version="99.0.0") as schema_file:

1134 config = self.make_instance(schema_file=schema_file)

1135 apdb = Apdb.from_config(config)

1136 self.assertEqual(

1137 apdb.schema.schemaVersion(),

1138 VersionTuple(99, 0, 0),

1139 )

1140

1141 def test_config_freeze(self) -> None:

1142 """Test that some config fields are correctly frozen in database."""

1143 config = self.make_instance()

1144

1145 # `enable_replica` is the only parameter that is frozen in all

1146 # implementations.

1147 config.enable_replica = not self.enable_replica

1148 apdb = Apdb.from_config(config)

1149 frozen_config = apdb.getConfig()

1150 self.assertEqual(frozen_config.enable_replica, self.enable_replica)

1151

1152

1153class ApdbSchemaUpdateTest(TestCaseMixin, ABC):

1154 """Base class for unit tests that verify how schema changes work."""

1155

1156 visit_time = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")

1157

1158 @abstractmethod

1159 def make_instance(self, **kwargs: Any) -> ApdbConfig:

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

1161

1162 This method should return configuration that point to the identical

1163 database instance on each call (i.e. ``db_url`` must be the same,

1164 which also means for sqlite it has to use on-disk storage).

1165 """

1166 raise NotImplementedError()

1167

1168 def make_region(self, xyz: tuple[float, float, float] = (1.0, 1.0, -1.0)) -> Region:

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

1170 return _make_region(xyz)

1171

1172 def test_schema_add_replica(self) -> None:

1173 """Check that new code can work with old schema without replica

1174 tables.

1175 """

1176 # Make schema without replica tables.

1177 config = self.make_instance(enable_replica=False)

1178 apdb = Apdb.from_config(config)

1179 apdb_replica = ApdbReplica.from_config(config)

1180

1181 # Make APDB instance configured for replication.

1182 config.enable_replica = True

1183 apdb = Apdb.from_config(config)

1184

1185 # Try to insert something, should work OK.

1186 region = self.make_region()

1187 visit_time = self.visit_time

1188

1189 # have to store Objects first

1190 objects = makeObjectCatalog(region, 100)

1191 sources = makeSourceCatalog(objects, visit_time)

1192 fsources = makeForcedSourceCatalog(objects, visit_time)

1193 apdb.store(visit_time, objects, sources, fsources)

1194

1195 # There should be no replica chunks.

1196 replica_chunks = apdb_replica.getReplicaChunks()

1197 self.assertIsNone(replica_chunks)

1198

1199 def test_schemaVersionCheck(self) -> None:

1200 """Check version number compatibility."""

1201 config = self.make_instance()

1202 apdb = Apdb.from_config(config)

1203

1204 self.assertEqual(apdb.schema.schemaVersion(), VersionTuple(0, 1, 1))

1205

1206 # Claim that schema version is now 99.0.0, must raise an exception.

1207 with update_schema_yaml(config.schema_file, version="99.0.0") as schema_file:

1208 config.schema_file = schema_file

1209 with self.assertRaises(IncompatibleVersionError):

1210 apdb = Apdb.from_config(config)

1211 # Version is checked only when we try to do connect.

1212 apdb.metadata.items()

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

601 statements