Coverage for tests/test_apdbSql.py: 48%

1# This file is part of dax_apdb. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <http://www.gnu.org/licenses/>. 

21 

22"""Unit test for Apdb class.""" 

23 

24import gc 

25import os 

26import shutil 

27import tempfile 

28import unittest 

29from typing import Any 

30from unittest.mock import patch 

31 

32import sqlalchemy 

33 

34import lsst.utils.tests 

35from lsst.dax.apdb import Apdb, ApdbConfig, ApdbReplica, ApdbTables, ApdbUpdateRecord, ReplicaChunk 

36from lsst.dax.apdb.pixelization import Pixelization 

37from lsst.dax.apdb.sql import ApdbSql, ApdbSqlConfig 

38from lsst.dax.apdb.tests import ApdbSchemaUpdateTest, ApdbTest 

39 

40try: 

41 import testing.postgresql 

42except ImportError: 

43 testing = None 

44 

45TEST_SCHEMA = os.path.join(os.path.abspath(os.path.dirname(__file__)), "config/schema.yaml") 

46# Schema that uses `datetime` for timestamps. 

47TEST_SCHEMA_DT = os.path.join(os.path.abspath(os.path.dirname(__file__)), "config/schema-datetime.yaml") 

48 

49 

50class ApdbSQLTest(ApdbTest): 

51 """A common base class for SQL APDB tests.""" 

52 

53 dia_object_index: str 

54 

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

56 """Create database and return its config.""" 

57 kw = { 

58 "schema_file": TEST_SCHEMA, 

59 "dia_object_index": self.dia_object_index, 

60 "enable_replica": self.enable_replica, 

61 } 

62 kw.update(kwargs) 

63 return ApdbSql.init_database(**kw) # type: ignore[arg-type] 

64 

65 def getDiaObjects_table(self) -> ApdbTables: 

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

67 return ApdbTables.DiaObject 

68 

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

70 """Return pixelization used by implementation.""" 

71 assert isinstance(config, ApdbSqlConfig), "Only expect ApdbSqlConfig here" 

72 return Pixelization("htm", config.pixelization.htm_level, config.pixelization.htm_max_ranges) 

73 

74 def test_connection_timeout(self) -> None: 

75 """Test that setting connection timeout does not break things.""" 

76 config = self.make_instance() 

77 assert isinstance(config, ApdbSqlConfig), "Only expect ApdbSqlConfig here" 

78 config.connection_config.connection_timeout = 60.0 

79 Apdb.from_config(config) 

80 

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

82 # Docstring inherited. 

83 assert isinstance(apdb, ApdbSql), "Expecting ApdbSql instance" 

84 apdb._storeUpdateRecords(records, chunk, store_chunk=True) 

85 

86 

87class ApdbSQLiteTestCase(ApdbSQLTest, unittest.TestCase): 

88 """A test case for ApdbSql class using SQLite backend.""" 

89 

90 fsrc_requires_id_list = True 

91 dia_object_index = "baseline" 

92 use_mjd = True 

93 

94 def setUp(self) -> None: 

95 self.tempdir = tempfile.mkdtemp() 

96 self.db_url = f"sqlite:///{self.tempdir}/apdb.sqlite3" 

97 

98 def tearDown(self) -> None: 

99 shutil.rmtree(self.tempdir, ignore_errors=True) 

100 

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

102 return super().make_instance(db_url=self.db_url, **kwargs) 

103 

104 

105class ApdbSQLiteDatetimeTestCase(ApdbSQLTest, unittest.TestCase): 

106 """A test case for ApdbSql class using SQLite backend and schema with 

107 timestamps in TAI MJD. 

108 """ 

109 

110 fsrc_requires_id_list = True 

111 dia_object_index = "baseline" 

112 use_mjd = False 

113 

114 def setUp(self) -> None: 

115 self.tempdir = tempfile.mkdtemp() 

116 self.db_url = f"sqlite:///{self.tempdir}/apdb.sqlite3" 

117 

118 def tearDown(self) -> None: 

119 shutil.rmtree(self.tempdir, ignore_errors=True) 

120 

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

122 if "schema_file" in kwargs: 

123 return super().make_instance(db_url=self.db_url, **kwargs) 

124 else: 

125 return super().make_instance(db_url=self.db_url, schema_file=TEST_SCHEMA_DT, **kwargs) 

126 

127 

128class ApdbSQLiteTestCaseLastObject(ApdbSQLiteTestCase): 

129 """A test case for ApdbSql class using SQLite backend and DiaObjectLast 

130 table. 

131 """ 

132 

133 dia_object_index = "last_object_table" 

134 

135 def getDiaObjects_table(self) -> ApdbTables: 

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

137 return ApdbTables.DiaObjectLast 

138 

139 

140class ApdbSQLiteTestCasePixIdIovIndex(ApdbSQLiteTestCase): 

141 """A test case for ApdbSql class using SQLite backend with pix_id_iov 

142 indexing. 

143 """ 

144 

145 dia_object_index = "pix_id_iov" 

146 

147 

148class ApdbSQLiteTestCaseReplica(ApdbSQLiteTestCase): 

149 """Test case for ApdbSql class using SQLite backend with replica tables.""" 

150 

151 enable_replica = True 

152 meta_row_count = 4 

153 

154 

155@unittest.skipUnless(testing is not None, "testing.postgresql module not found") 

156class ApdbPostgresTestCase(ApdbSQLTest, unittest.TestCase): 

157 """A test case for ApdbSql class using Postgres backend.""" 

158 

159 fsrc_requires_id_list = True 

160 dia_object_index = "last_object_table" 

161 postgresql: Any 

162 enable_replica = True 

163 meta_row_count = 4 

164 timestamp_type_name = "datetime64[ns]" 

165 

166 @classmethod 

167 def setUpClass(cls) -> None: 

168 # Create the postgres test server. 

169 cls.postgresql = testing.postgresql.PostgresqlFactory(cache_initialized_db=True) 

170 super().setUpClass() 

171 

172 @classmethod 

173 def tearDownClass(cls) -> None: 

174 # Clean up any lingering SQLAlchemy engines/connections 

175 # so they're closed before we shut down the server. 

176 gc.collect() 

177 cls.postgresql.clear_cache() 

178 super().tearDownClass() 

179 

180 def setUp(self) -> None: 

181 self.server = self.postgresql() 

182 

183 def tearDown(self) -> None: 

184 self.server = self.postgresql() 

185 

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

187 return super().make_instance(db_url=self.server.url(), **kwargs) 

188 

189 def getDiaObjects_table(self) -> ApdbTables: 

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

191 return ApdbTables.DiaObjectLast 

192 

193 

194@unittest.skipUnless(testing is not None, "testing.postgresql module not found") 

195class ApdbPostgresNamespaceTestCase(ApdbPostgresTestCase): 

196 """A test case for ApdbSql class using Postgres backend with schema name""" 

197 

198 # use mixed case to trigger quoting 

199 namespace = "ApdbSchema" 

200 

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

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

203 return super().make_instance(namespace=self.namespace, **kwargs) 

204 

205 

206class ApdbSchemaUpdateSQLiteTestCase(ApdbSchemaUpdateTest, unittest.TestCase): 

207 """A test case for schema updates using SQLite backend.""" 

208 

209 def setUp(self) -> None: 

210 self.tempdir = tempfile.mkdtemp() 

211 self.db_url = f"sqlite:///{self.tempdir}/apdb.sqlite3" 

212 

213 def tearDown(self) -> None: 

214 shutil.rmtree(self.tempdir, ignore_errors=True) 

215 

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

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

218 kw = { 

219 "db_url": self.db_url, 

220 "schema_file": TEST_SCHEMA, 

221 } 

222 kw.update(kwargs) 

223 return ApdbSql.init_database(**kw) # type: ignore[arg-type] 

224 

225 

226class ApdbSQLiteFromUriTestCase(unittest.TestCase): 

227 """A test case for for instantiating ApdbSql via URI.""" 

228 

229 def setUp(self) -> None: 

230 self.tempdir = tempfile.mkdtemp() 

231 self.addCleanup(shutil.rmtree, self.tempdir, ignore_errors=True) 

232 self.db_url = f"sqlite:///{self.tempdir}/apdb.sqlite3" 

233 config = ApdbSql.init_database(db_url=self.db_url, schema_file=TEST_SCHEMA) 

234 # TODO: This will need update when we switch to pydantic configs. 

235 self.config_path = os.path.join(self.tempdir, "apdb-config.yaml") 

236 config.save(self.config_path) 

237 self.bad_config_path = os.path.join(self.tempdir, "not-config.yaml") 

238 self.index_path = os.path.join(self.tempdir, "apdb-index.yaml") 

239 with open(self.index_path, "w") as index_file: 

240 print(f'label1: "{self.config_path}"', file=index_file) 

241 print(f'"label2/pex_config": "{self.config_path}"', file=index_file) 

242 print(f'bad-label: "{self.bad_config_path}"', file=index_file) 

243 # File with incorrect format. 

244 self.bad_index_path = os.path.join(self.tempdir, "apdb-index-bad.yaml") 

245 with open(self.bad_index_path, "w") as index_file: 

246 print(f'label1: ["{self.config_path}"]', file=index_file) 

247 self.missing_index_path = os.path.join(self.tempdir, "no-apdb-index.yaml") 

248 

249 def test_make_apdb_from_path(self) -> None: 

250 """Check that we can make APDB instance from config URI.""" 

251 Apdb.from_uri(self.config_path) 

252 with self.assertRaises(FileNotFoundError): 

253 Apdb.from_uri(self.bad_config_path) 

254 

255 def test_make_apdb_from_labels(self) -> None: 

256 """Check that we can make APDB instance from config URI.""" 

257 # Replace DAX_APDB_INDEX_URI value 

258 new_env = {"DAX_APDB_INDEX_URI": self.index_path} 

259 with patch.dict(os.environ, new_env, clear=True): 

260 Apdb.from_uri("label:label1") 

261 Apdb.from_uri("label:label2") 

262 # Label does not exist. 

263 with self.assertRaises(ValueError): 

264 Apdb.from_uri("label:not-a-label") 

265 # Label exists but points to a missing config. 

266 with self.assertRaises(FileNotFoundError): 

267 Apdb.from_uri("label:bad-label") 

268 

269 def test_make_apdb_bad_index(self) -> None: 

270 """Check what happens when DAX_APDB_INDEX_URI is broken.""" 

271 # envvar is set but empty. 

272 new_env = {"DAX_APDB_INDEX_URI": ""} 

273 with patch.dict(os.environ, new_env, clear=True): 

274 with self.assertRaises(RuntimeError): 

275 Apdb.from_uri("label:label") 

276 

277 # envvar is set to something non-existing. 

278 new_env = {"DAX_APDB_INDEX_URI": self.missing_index_path} 

279 with patch.dict(os.environ, new_env, clear=True): 

280 with self.assertRaises(FileNotFoundError): 

281 Apdb.from_uri("label:label") 

282 

283 # envvar points to an incorrect file. 

284 new_env = {"DAX_APDB_INDEX_URI": self.bad_index_path} 

285 with patch.dict(os.environ, new_env, clear=True): 

286 with self.assertRaises(TypeError): 

287 Apdb.from_uri("label:label") 

288 

289 def test_remove_database_file(self) -> None: 

290 """Check that SQLite does not try to recreate database after it was 

291 removed, but config persisted. 

292 """ 

293 os.unlink(f"{self.tempdir}/apdb.sqlite3") 

294 with self.assertRaisesRegex(sqlalchemy.exc.OperationalError, "unable to open database file"): 

295 Apdb.from_uri(self.config_path) 

296 

297 def test_make_apdb_replica(self) -> None: 

298 """Check that we can make ApdbReplica instance from config URI.""" 

299 ApdbReplica.from_uri(self.config_path) 

300 with self.assertRaises(FileNotFoundError): 

301 Apdb.from_uri(self.bad_config_path) 

302 

303 

304class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase): 

305 """Run file leak tests.""" 

306 

307 

308def setup_module(module: Any) -> None: 

309 """Configure pytest.""" 

310 lsst.utils.tests.init() 

311 

312 

313if __name__ == "__main__": 

314 lsst.utils.tests.init() 

315 unittest.main()