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_null_integer_type(self) -> None:

471 """Test that integer column with NULLs correct type on select."""

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 # have to store Objects first

479 objects = makeObjectCatalog(region, 100)

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

481 # Reset some diaObjectIds to NULL.

482 sources.loc[0:10, "diaObjectId"] = None

483

484 # save the objects and sources

485 apdb.store(visit_time, objects, sources)

486

487 # read it back, no ID filtering

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

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

490 assert res is not None, "Expecting catalog, not None"

491 self.assertEqual(res.dtypes["diaObjectId"], pandas.Int64Dtype())

492

493 def test_timestamps(self) -> None:

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

495 config = self.make_instance()

496 apdb = Apdb.from_config(config)

497

498 region = self.make_region()

499 visit_time = self.visit_time

500

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

502 # truncated returned values.

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

504 objects = makeObjectCatalog(region, 100)

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

506 catalog = makeForcedSourceCatalog(

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

508 )

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

510

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

512

513 # read it back and check sizes

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

515 assert res is not None

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

517

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

519 self.assertIn(time_processed_column, res.dtypes)

520 dtype = res.dtypes[time_processed_column]

521 timestamp_type_names = (

522 ("float64",) if self.use_mjd else ("datetime64[ms]", "datetime64[us]", "datetime64[ns]")

523 )

524 self.assertIn(dtype.name, timestamp_type_names)

525 # Verify that returned time is sensible.

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

527

528 def test_getDiaObjectsForDedup(self) -> None:

529 """Test getDiaObjectsForDedup() method."""

530 config = self.make_instance()

531 apdb = Apdb.from_config(config)

532

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

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

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

536 nobj = 100

537 objects1 = makeObjectCatalog(region1, nobj)

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

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

540

541 visits = [

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

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

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

545 ]

546

547 for visit_time, objects in visits:

548 apdb.store(visit_time, objects)

549

550 catalog = apdb.getDiaObjectsForDedup()

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

552

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

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

555

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

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

558

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

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

561

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

563 catalog = apdb.getDiaObjectsForDedup(time)

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

565

566 def test_getDiaSourcesForDiaObjects(self) -> None:

567 """Test getDiaSourcesForDiaObjects() method."""

568 config = self.make_instance()

569 apdb = Apdb.from_config(config)

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

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

572

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

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

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

576 nobj = 100

577 objects1 = makeObjectCatalog(region1, nobj)

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

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

580

581 visits = [

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

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

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

585 ]

586

587 start_id = 1_000_000

588 for visit_time, objects in visits:

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

590 apdb.store(visit_time, objects, sources)

591 start_id += 1_000_000

592

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

594 object_ids = [

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

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

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

598 ]

599

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

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

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

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

604

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

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

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

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

609

610 def test_reassignDiaSourcesToDiaObjects(self) -> None:

611 """Test reassignDiaSourcesToDiaObjects() method."""

612 config = self.make_instance()

613 apdb = Apdb.from_config(config)

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

615 apdb_replica = ApdbReplica.from_config(config)

616

617 visit_time = self.visit_time

618 lonlat1 = LonLat.fromDegrees(0.0, 0.0)

619 lonlat2 = LonLat.fromDegrees(180.0, 0.0)

620 # regons around lonlat1/2

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

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

623

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

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

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

627 apdb.store(visit_time, objects, sources)

628

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

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

631 apdb.store(visit_time, objects, sources)

632

633 # check that everything as we think it is.

634 objects = apdb.getDiaObjects(region1)

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

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

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

638 assert sources is not None

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

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

641

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

643 midpointMjdTai = visit_time.tai.mjd

644 reassign = {

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

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

647 }

648 apdb.reassignDiaSourcesToDiaObjects(reassign)

649

650 objects = apdb.getDiaObjects(region1)

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

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

653 assert sources is not None

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

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

656

657 # Reassign but do not increment/decrement nDiaSources.

658 reassign = {

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

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

661 }

662 apdb.reassignDiaSourcesToDiaObjects(

663 reassign, increment_nDiaSources=False, decrement_nDiaSources=False

664 )

665

666 objects = apdb.getDiaObjects(region2)

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

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

669 assert sources is not None

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

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

672

673 replica_chunks = apdb_replica.getReplicaChunks()

674 if not self.enable_replica:

675 self.assertIsNone(replica_chunks)

676 else:

677 assert replica_chunks is not None

678

679 # There could be one or two chunks.

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

681

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

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

684 # that region, plus two reassignments for region2 without

685 # increments/decrements.

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

687

688 def test_setValidityEnd(self) -> None:

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

690 # don't care about sources.

691 config = self.make_instance()

692 apdb = Apdb.from_config(config)

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

694 apdb_replica = ApdbReplica.from_config(config)

695

696 region = self.make_region()

697 visit_time = self.visit_time

698

699 # make catalog with Objects

700 catalog = makeObjectCatalog(region, 100)

701

702 # store catalog

703 apdb.store(visit_time, catalog)

704

705 # read it back and check sizes

706 res = apdb.getDiaObjects(region)

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

708

709 # Select first 10 objects.

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

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

712 self.assertEqual(count, 10)

713

714 res = apdb.getDiaObjects(region)

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

716

717 replica_chunks = apdb_replica.getReplicaChunks()

718 if not self.enable_replica:

719 self.assertIsNone(replica_chunks)

720 else:

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

722 assert replica_chunks is not None

723

724 # There could be one or two chunks.

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

726

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

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

729

730 # Check that empty list works.

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

732 self.assertEqual(count, 0)

733

734 # Try with non-existing object.

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

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

737 with self.assertRaises(LookupError):

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

739

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

741 self.assertEqual(count, 2)

742

743 def test_resetDedup(self) -> None:

744 """Test resetDedup method."""

745 # don't care about sources.

746 config = self.make_instance()

747 apdb = Apdb.from_config(config)

748

749 region = self.make_region()

750

751 # make catalog with Objects

752 objects = makeObjectCatalog(region, 100)

753

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

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

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

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

758

759 # store catalog

760 apdb.store(visit_time1, objects)

761

762 catalog = apdb.getDiaObjectsForDedup()

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

764

765 catalog = apdb.getDiaObjectsForDedup(visit_time1)

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

767

768 apdb.resetDedup(dedup_time1)

769

770 catalog = apdb.getDiaObjectsForDedup(visit_time1)

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

772

773 apdb.store(visit_time2, objects)

774

775 catalog = apdb.getDiaObjectsForDedup()

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

777

778 catalog = apdb.getDiaObjectsForDedup(dedup_time1)

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

780

781 apdb.resetDedup(dedup_time2)

782

783 catalog = apdb.getDiaObjectsForDedup(dedup_time1)

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

785

786 catalog = apdb.getDiaObjectsForDedup()

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

788

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

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

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

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

793 by resetDedup().

794 """

795 raise NotImplementedError()

796

797 def test_getChunks(self) -> None:

798 """Store and retrieve replica chunks."""

799 # don't care about sources.

800 config = self.make_instance()

801 apdb = Apdb.from_config(config)

802 apdb_replica = ApdbReplica.from_config(config)

803 visit_time = self.visit_time

804

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

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

807 nobj = 100

808 objects1 = makeObjectCatalog(region1, nobj)

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

810

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

812 # records.

813 visits = [

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

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

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

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

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

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

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

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

822 ]

823

824 start_id = 0

825 for visit_time, objects in visits:

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

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

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

829 start_id += nobj

830

831 replica_chunks = apdb_replica.getReplicaChunks()

832 if not self.enable_replica:

833 self.assertIsNone(replica_chunks)

834

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

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

837

838 else:

839 assert replica_chunks is not None

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

841

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

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

844

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

846 if n_records is None:

847 n_records = len(replica_chunks) * nobj

848 res = apdb_replica.getTableDataChunks(

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

850 )

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

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

853 validityStartType = (

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

855 )

856 self.assert_column_types(

857 res,

858 {

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

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

861 validityStartColumn: validityStartType,

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

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

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

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

866 },

867 )

868

869 res = apdb_replica.getTableDataChunks(

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

871 )

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

873 self.assert_column_types(

874 res,

875 {

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

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

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

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

880 },

881 )

882

883 res = apdb_replica.getTableDataChunks(

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

885 )

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

887 self.assert_column_types(

888 res,

889 {

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

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

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

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

894 },

895 )

896

897 # read it back and check sizes

898 _check_chunks(replica_chunks, 800)

899 _check_chunks(replica_chunks[1:], 600)

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

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

902 _check_chunks([])

903

904 # try to remove some of those

905 deleted_chunks = replica_chunks[:1]

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

907

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

909 _check_chunks(deleted_chunks, 0)

910

911 replica_chunks = apdb_replica.getReplicaChunks()

912 assert replica_chunks is not None

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

914

915 _check_chunks(replica_chunks, 600)

916

917 def test_reassignObjects(self) -> None:

918 """Reassign DiaObjects."""

919 # don't care about sources.

920 config = self.make_instance()

921 apdb = Apdb.from_config(config)

922

923 region = self.make_region()

924 visit_time = self.visit_time

925 objects = makeObjectCatalog(region, 100)

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

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

928 apdb.store(visit_time, objects, sources)

929

930 # read it back and filter by ID

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

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

933

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

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

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

937

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

939 apdb.reassignDiaSources(

940 {

941 1000: 1,

942 7: 3,

943 }

944 )

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

946

947 def test_storeUpdateRecord(self) -> None:

948 """Test _storeUpdateRecord() method."""

949 config = self.make_instance()

950 apdb = Apdb.from_config(config)

951

952 # Times are totally arbitrary.

953 update_time_ns1 = 2_000_000_000_000_000_000

954 update_time_ns2 = 2_000_000_001_000_000_000

955 records = [

956 ApdbReassignDiaSourceToSSObjectRecord(

957 update_time_ns=update_time_ns1,

958 update_order=0,

959 diaSourceId=1,

960 ssObjectId=1,

961 ssObjectReassocTimeMjdTai=60000.0,

962 ra=45.0,

963 dec=-45.0,

964 midpointMjdTai=60000.0,

965 ),

966 ApdbWithdrawDiaSourceRecord(

967 update_time_ns=update_time_ns1,

968 update_order=1,

969 diaSourceId=123456,

970 timeWithdrawnMjdTai=61000.0,

971 ra=45.0,

972 dec=-45.0,

973 midpointMjdTai=60000.0,

974 ),

975 ApdbReassignDiaSourceToSSObjectRecord(

976 update_time_ns=update_time_ns1,

977 update_order=3,

978 diaSourceId=2,

979 ssObjectId=3,

980 ssObjectReassocTimeMjdTai=60000.0,

981 ra=45.0,

982 dec=-45.0,

983 midpointMjdTai=60000.0,

984 ),

985 ApdbWithdrawDiaSourceRecord(

986 update_time_ns=update_time_ns2,

987 update_order=0,

988 diaSourceId=123456,

989 timeWithdrawnMjdTai=61000.0,

990 ra=45.0,

991 dec=-45.0,

992 midpointMjdTai=60000.0,

993 ),

994 ]

995

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

997 chunk = ReplicaChunk.make_replica_chunk(update_time, 600)

998

999 if not self.enable_replica:

1000 with self.assertRaises(TypeError):

1001 self.store_update_records(apdb, records, chunk)

1002 else:

1003 self.store_update_records(apdb, records, chunk)

1004

1005 apdb_replica = ApdbReplica.from_config(config)

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

1007

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

1009 self.assertEqual(records_returned, records)

1010

1011 @abstractmethod

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

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

1014 raise NotImplementedError()

1015

1016 def test_midpointMjdTai_src(self) -> None:

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

1018 config = self.make_instance()

1019 apdb = Apdb.from_config(config)

1020

1021 region = self.make_region()

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

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

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

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

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

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

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

1029

1030 objects = makeObjectCatalog(region, 100)

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

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

1033 apdb.store(src_time1, objects, sources)

1034

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

1036 apdb.store(src_time2, objects, sources)

1037

1038 # reading at time of last save should read all

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

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

1041

1042 # one second before 12 months

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

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

1045

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

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

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

1049

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

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

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

1053

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

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

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

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

1058

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

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

1061

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

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

1064

1065 def test_midpointMjdTai_fsrc(self) -> None:

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

1067 config = self.make_instance()

1068 apdb = Apdb.from_config(config)

1069

1070 region = self.make_region()

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

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

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

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

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

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

1077

1078 objects = makeObjectCatalog(region, 100)

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

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

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

1082

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

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

1085

1086 # reading at time of last save should read all

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

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

1089

1090 # one second before 12 months

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

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

1093

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

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

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

1097

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

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

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

1101

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

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

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

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

1106

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

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

1109

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

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

1112

1113 def test_metadata(self) -> None:

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

1115 config = self.make_instance()

1116 apdb = Apdb.from_config(config)

1117 metadata = apdb.metadata

1118

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

1120 # and a frozen JSON config.

1121 self.assertFalse(metadata.empty())

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

1123

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

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

1126

1127 self.assertFalse(metadata.empty())

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

1129

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

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

1132

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

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

1135

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

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

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

1139

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

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

1142

1143 def test_schemaVersionFromYaml(self) -> None:

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

1145 config = self.make_instance()

1146 default_schema = config.schema_file

1147 apdb = Apdb.from_config(config)

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

1149

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

1151 config = self.make_instance(schema_file=schema_file)

1152 apdb = Apdb.from_config(config)

1153 self.assertEqual(

1154 apdb.schema.schemaVersion(),

1155 VersionTuple(0, 1, 0),

1156 )

1157

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

1159 config = self.make_instance(schema_file=schema_file)

1160 apdb = Apdb.from_config(config)

1161 self.assertEqual(

1162 apdb.schema.schemaVersion(),

1163 VersionTuple(99, 0, 0),

1164 )

1165

1166 def test_config_freeze(self) -> None:

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

1168 config = self.make_instance()

1169

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

1171 # implementations.

1172 config.enable_replica = not self.enable_replica

1173 apdb = Apdb.from_config(config)

1174 frozen_config = apdb.getConfig()

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

1176

1177

1178class ApdbSchemaUpdateTest(TestCaseMixin, ABC):

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

1180

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

1182

1183 @abstractmethod

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

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

1186

1187 This method should return configuration that point to the identical

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

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

1190 """

1191 raise NotImplementedError()

1192

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

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

1195 return _make_region(xyz)

1196

1197 def test_schema_add_replica(self) -> None:

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

1199 tables.

1200 """

1201 # Make schema without replica tables.

1202 config = self.make_instance(enable_replica=False)

1203 apdb = Apdb.from_config(config)

1204 apdb_replica = ApdbReplica.from_config(config)

1205

1206 # Make APDB instance configured for replication.

1207 config.enable_replica = True

1208 apdb = Apdb.from_config(config)

1209

1210 # Try to insert something, should work OK.

1211 region = self.make_region()

1212 visit_time = self.visit_time

1213

1214 # have to store Objects first

1215 objects = makeObjectCatalog(region, 100)

1216 sources = makeSourceCatalog(objects, visit_time)

1217 fsources = makeForcedSourceCatalog(objects, visit_time)

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

1219

1220 # There should be no replica chunks.

1221 replica_chunks = apdb_replica.getReplicaChunks()

1222 self.assertIsNone(replica_chunks)

1223

1224 def test_schemaVersionCheck(self) -> None:

1225 """Check version number compatibility."""

1226 config = self.make_instance()

1227 apdb = Apdb.from_config(config)

1228

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

1230

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

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

1233 config.schema_file = schema_file

1234 with self.assertRaises(IncompatibleVersionError):

1235 apdb = Apdb.from_config(config)

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

1237 apdb.metadata.items()

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

614 statements