Coverage for tests/test_tap_schema.py: 19%

1# This file is part of felis. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://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 <https://www.gnu.org/licenses/>. 

21 

22import os 

23import shutil 

24import tempfile 

25import unittest 

26from typing import Any 

27 

28from sqlalchemy import create_engine, select 

29 

30from felis.datamodel import Schema 

31from felis.tap_schema import DataLoader, TableManager 

32 

33TEST_DIR = os.path.dirname(__file__) 

34TEST_SALES = os.path.join(TEST_DIR, "data", "sales.yaml") 

35TEST_TAP_SCHEMA = os.path.join(TEST_DIR, "data", "test_tap_schema.yaml") 

36TEST_COMPOSITE_KEYS = os.path.join(TEST_DIR, "data", "test_composite_keys.yaml") 

37 

38 

39class TableManagerTestCase(unittest.TestCase): 

40 """Test the `TableManager` class.""" 

41 

42 def setUp(self) -> None: 

43 """Set up the test case.""" 

44 with open(TEST_SALES) as test_file: 

45 self.schema = Schema.from_stream(test_file) 

46 

47 def test_create_table_manager(self) -> None: 

48 """Test the TAP table manager class.""" 

49 mgr = TableManager() 

50 

51 schema_name = mgr.schema_name 

52 

53 # Check the created metadata and tables. 

54 self.assertNotEqual(len(mgr.metadata.tables), 0) 

55 self.assertEqual(mgr.metadata.schema, schema_name) 

56 for table_name in mgr.get_table_names_std(): 

57 mgr[table_name] 

58 

59 # Make sure that creating a new table manager works when one has 

60 # already been created. 

61 mgr = TableManager() 

62 

63 def test_table_name_postfix(self) -> None: 

64 """Test the table name postfix.""" 

65 mgr = TableManager(apply_schema_to_metadata=False, table_name_postfix="_test") 

66 for table_name in mgr.metadata.tables: 

67 self.assertTrue(table_name.endswith("_test")) 

68 

69 

70class DataLoaderTestCase(unittest.TestCase): 

71 """Test the `DataLoader` class.""" 

72 

73 def setUp(self) -> None: 

74 """Set up the test case.""" 

75 with open(TEST_TAP_SCHEMA) as test_file: 

76 self.schema = Schema.from_stream(test_file, context={"id_generation": True}) 

77 

78 self.tmpdir = tempfile.mkdtemp(dir=TEST_DIR) 

79 

80 def tearDown(self) -> None: 

81 """Clean up temporary directory.""" 

82 shutil.rmtree(self.tmpdir, ignore_errors=True) 

83 

84 def test_sqlite(self) -> None: 

85 """Test the `DataLoader` using an in-memory SQLite database.""" 

86 engine = create_engine("sqlite:///:memory:") 

87 

88 mgr = TableManager(apply_schema_to_metadata=False) 

89 mgr.initialize_database(engine) 

90 

91 loader = DataLoader(self.schema, mgr, engine) 

92 loader.load() 

93 

94 def test_sql_output(self) -> None: 

95 """Test printing SQL to stdout and writing SQL to a file.""" 

96 engine = create_engine("sqlite:///:memory:") 

97 mgr = TableManager(apply_schema_to_metadata=False) 

98 loader = DataLoader(self.schema, mgr, engine, dry_run=True, print_sql=True) 

99 loader.load() 

100 

101 sql_path = os.path.join(self.tmpdir, "test_tap_schema_print_sql.sql") 

102 loader = DataLoader(self.schema, mgr, engine, dry_run=True, print_sql=True, output_path=sql_path) 

103 loader.load() 

104 

105 self.assertTrue(os.path.exists(sql_path)) 

106 with open(sql_path) as sql_file: 

107 sql_data = sql_file.read() 

108 insert_count = sql_data.count("INSERT INTO") 

109 self.assertEqual( 

110 insert_count, 

111 22, 

112 f"Expected 22 'INSERT INTO' statements, found {insert_count}", 

113 ) 

114 

115 def test_unique_keys(self) -> None: 

116 """Test generation of unique foreign keys.""" 

117 engine = create_engine("sqlite:///:memory:") 

118 

119 mgr = TableManager(apply_schema_to_metadata=False) 

120 mgr.initialize_database(engine) 

121 

122 loader = DataLoader(self.schema, mgr, engine, unique_keys=True) 

123 loader.load() 

124 

125 keys_data = mgr.select(engine, "keys") 

126 self.assertGreaterEqual(len(keys_data), 1) 

127 for row in keys_data: 

128 self.assertTrue(row["key_id"].startswith(f"{self.schema.name}_")) 

129 

130 key_columns_data = mgr.select(engine, "key_columns") 

131 self.assertGreaterEqual(len(key_columns_data), 1) 

132 for row in key_columns_data: 

133 self.assertTrue(row["key_id"].startswith(f"{self.schema.name}_")) 

134 

135 def test_select_with_filter(self) -> None: 

136 """Test selecting rows with a filter.""" 

137 engine = create_engine("sqlite:///:memory:") 

138 mgr = TableManager(apply_schema_to_metadata=False) 

139 mgr.initialize_database(engine) 

140 loader = DataLoader(self.schema, mgr, engine, unique_keys=True) 

141 loader.load() 

142 

143 rows = mgr.select(engine, "columns", "table_name = 'test_schema.table1'") 

144 self.assertEqual(len(rows), 16) 

145 

146 

147def _find_row(rows: list[dict[str, Any]], column_name: str, value: str) -> dict[str, Any]: 

148 next_row = next( 

149 (row for row in rows if row[column_name] == value), 

150 None, 

151 ) 

152 assert next_row is not None 

153 assert isinstance(next_row, dict) 

154 return next_row 

155 

156 

157class TapSchemaSqliteSetup: 

158 """Set up the TAP_SCHEMA SQLite database for testing. 

159 

160 Parameters 

161 ---------- 

162 test_file_path: 

163 Path to the TAP_SCHEMA test file. 

164 

165 context 

166 Context for the schema. Default is an empty dictionary. 

167 """ 

168 

169 def __init__(self, test_file_path: str, context: dict = {}) -> None: 

170 with open(test_file_path) as test_file: 

171 self._schema = Schema.from_stream(test_file, context=context) 

172 

173 self._engine = create_engine("sqlite:///:memory:") 

174 

175 mgr = TableManager(apply_schema_to_metadata=False) 

176 mgr.initialize_database(self._engine) 

177 self._mgr = mgr 

178 

179 loader = DataLoader(self._schema, mgr, self._engine, tap_schema_index=2) 

180 loader.load() 

181 

182 @property 

183 def schema(self) -> Schema: 

184 """Return the schema.""" 

185 return self._schema 

186 

187 @property 

188 def engine(self) -> Any: 

189 """Return the engine.""" 

190 return self._engine 

191 

192 @property 

193 def mgr(self) -> TableManager: 

194 """Return the table manager.""" 

195 return self._mgr 

196 

197 

198class TapSchemaDataTest(unittest.TestCase): 

199 """Test the validity of generated TAP SCHEMA data.""" 

200 

201 def setUp(self) -> None: 

202 """Set up the test case.""" 

203 self.tap_schema_setup = TapSchemaSqliteSetup(TEST_TAP_SCHEMA, context={"id_generation": True}) 

204 

205 def test_schemas(self) -> None: 

206 schemas_table = self.tap_schema_setup.mgr["schemas"] 

207 with self.tap_schema_setup.engine.connect() as connection: 

208 result = connection.execute(select(schemas_table)) 

209 schema_data = [row._asdict() for row in result] 

210 

211 self.assertEqual(len(schema_data), 1) 

212 

213 schema = schema_data[0] 

214 self.assertEqual(schema["schema_name"], "test_schema") 

215 self.assertEqual(schema["description"], "Test schema") 

216 self.assertEqual(schema["utype"], "Schema") 

217 self.assertEqual(schema["schema_index"], 2) 

218 

219 def test_tables(self) -> None: 

220 tables_table = self.tap_schema_setup.mgr["tables"] 

221 with self.tap_schema_setup.engine.connect() as connection: 

222 result = connection.execute(select(tables_table)) 

223 table_data = [row._asdict() for row in result] 

224 

225 self.assertEqual(len(table_data), 2) 

226 

227 table = table_data[0] 

228 assert isinstance(table, dict) 

229 self.assertEqual(table["schema_name"], "test_schema") 

230 self.assertEqual(table["table_name"], f"{self.tap_schema_setup.schema.name}.table1") 

231 self.assertEqual(table["table_type"], "table") 

232 self.assertEqual(table["utype"], "Table") 

233 self.assertEqual(table["description"], "Test table 1") 

234 self.assertEqual(table["table_index"], 2) 

235 

236 def test_columns(self) -> None: 

237 columns_table = self.tap_schema_setup.mgr["columns"] 

238 with self.tap_schema_setup.engine.connect() as connection: 

239 result = connection.execute(select(columns_table)) 

240 column_data = [row._asdict() for row in result] 

241 

242 table1_rows = [ 

243 row for row in column_data if row["table_name"] == f"{self.tap_schema_setup.schema.name}.table1" 

244 ] 

245 self.assertNotEqual(len(table1_rows), 0) 

246 

247 boolean_col = _find_row(table1_rows, "column_name", "boolean_field") 

248 self.assertEqual(boolean_col["datatype"], "boolean") 

249 self.assertEqual(boolean_col["arraysize"], None) 

250 

251 byte_col = _find_row(table1_rows, "column_name", "byte_field") 

252 self.assertEqual(byte_col["datatype"], "unsignedByte") 

253 self.assertEqual(byte_col["arraysize"], None) 

254 

255 short_col = _find_row(table1_rows, "column_name", "short_field") 

256 self.assertEqual(short_col["datatype"], "short") 

257 self.assertEqual(short_col["arraysize"], None) 

258 

259 int_col = _find_row(table1_rows, "column_name", "int_field") 

260 self.assertEqual(int_col["datatype"], "int") 

261 self.assertEqual(int_col["arraysize"], None) 

262 

263 float_col = _find_row(table1_rows, "column_name", "float_field") 

264 self.assertEqual(float_col["datatype"], "float") 

265 self.assertEqual(float_col["arraysize"], None) 

266 

267 double_col = _find_row(table1_rows, "column_name", "double_field") 

268 self.assertEqual(double_col["datatype"], "double") 

269 self.assertEqual(double_col["arraysize"], None) 

270 

271 long_col = _find_row(table1_rows, "column_name", "long_field") 

272 self.assertEqual(long_col["datatype"], "long") 

273 self.assertEqual(long_col["arraysize"], None) 

274 

275 unicode_col = _find_row(table1_rows, "column_name", "unicode_field") 

276 self.assertEqual(unicode_col["datatype"], "unicodeChar") 

277 self.assertEqual(unicode_col["arraysize"], "128*") 

278 

279 binary_col = _find_row(table1_rows, "column_name", "binary_field") 

280 self.assertEqual(binary_col["datatype"], "unsignedByte") 

281 self.assertEqual(binary_col["arraysize"], "1024*") 

282 

283 ts = _find_row(table1_rows, "column_name", "timestamp_field") 

284 self.assertEqual(ts["datatype"], "char") 

285 self.assertEqual(ts["xtype"], "timestamp") 

286 self.assertEqual(ts["description"], "Timestamp field") 

287 self.assertEqual(ts["utype"], "Obs:Timestamp") 

288 self.assertEqual(ts["unit"], "s") 

289 self.assertEqual(ts["ucd"], "time.epoch") 

290 self.assertEqual(ts["principal"], 1) 

291 self.assertEqual(ts["std"], 1) 

292 self.assertEqual(ts["column_index"], 42) 

293 self.assertEqual(ts["size"], None) 

294 self.assertEqual(ts["arraysize"], "*") 

295 

296 char_col = _find_row(table1_rows, "column_name", "char_field") 

297 self.assertEqual(char_col["datatype"], "char") 

298 self.assertEqual(char_col["arraysize"], "64") 

299 

300 str_col = _find_row(table1_rows, "column_name", "string_field") 

301 self.assertEqual(str_col["datatype"], "char") 

302 self.assertEqual(str_col["arraysize"], "256*") 

303 

304 txt_col = _find_row(table1_rows, "column_name", "text_field") 

305 self.assertEqual(txt_col["datatype"], "char") 

306 self.assertEqual(txt_col["arraysize"], "*") 

307 

308 def test_keys(self) -> None: 

309 keys_table = self.tap_schema_setup.mgr["keys"] 

310 with self.tap_schema_setup.engine.connect() as connection: 

311 result = connection.execute(select(keys_table)) 

312 key_data = [row._asdict() for row in result] 

313 

314 self.assertEqual(len(key_data), 1) 

315 

316 key = key_data[0] 

317 assert isinstance(key, dict) 

318 

319 self.assertEqual(key["key_id"], "fk_table1_to_table2") 

320 self.assertEqual(key["from_table"], f"{self.tap_schema_setup.schema.name}.table1") 

321 self.assertEqual(key["target_table"], f"{self.tap_schema_setup.schema.name}.table2") 

322 self.assertEqual(key["description"], "Foreign key from table1 to table2") 

323 self.assertEqual(key["utype"], "ForeignKey") 

324 

325 def test_key_columns(self) -> None: 

326 key_columns_table = self.tap_schema_setup.mgr["key_columns"] 

327 with self.tap_schema_setup.engine.connect() as connection: 

328 result = connection.execute(select(key_columns_table)) 

329 key_column_data = [row._asdict() for row in result] 

330 

331 self.assertEqual(len(key_column_data), 1) 

332 

333 key_column = key_column_data[0] 

334 assert isinstance(key_column, dict) 

335 

336 self.assertEqual(key_column["key_id"], "fk_table1_to_table2") 

337 self.assertEqual(key_column["from_column"], "fk") 

338 self.assertEqual(key_column["target_column"], "id") 

339 

340 def test_bad_table_name(self) -> None: 

341 """Test getting a bad TAP_SCHEMA table name.""" 

342 with self.assertRaises(KeyError): 

343 self.tap_schema_setup.mgr["bad_table"] 

344 

345 

346class ForceUnboundArraySizeTest(unittest.TestCase): 

347 """Test that arraysize for appropriate types is set to '*' when the 

348 ``force_unboundeded_arraysize`` context flag is set to ``True``. 

349 """ 

350 

351 def setUp(self) -> None: 

352 """Set up the test case.""" 

353 self.tap_schema_setup = TapSchemaSqliteSetup( 

354 TEST_TAP_SCHEMA, context={"id_generation": True, "force_unbounded_arraysize": True} 

355 ) 

356 

357 def test_force_unbounded_arraysize(self) -> None: 

358 """Test that unbounded arraysize is set to None.""" 

359 columns_table = self.tap_schema_setup.mgr["columns"] 

360 with self.tap_schema_setup.engine.connect() as connection: 

361 result = connection.execute(select(columns_table)) 

362 column_data = [row._asdict() for row in result] 

363 

364 table1_rows = [ 

365 row for row in column_data if row["table_name"] == f"{self.tap_schema_setup.schema.name}.table1" 

366 ] 

367 for row in table1_rows: 

368 if row["column_name"] in ["string_field", "text_field", "unicode_field", "binary_field"]: 

369 self.assertEqual(row["arraysize"], "*") 

370 

371 

372class CompositeKeysTestCase(unittest.TestCase): 

373 """Test the handling of composite foreign keys.""" 

374 

375 def setUp(self) -> None: 

376 """Set up the test case.""" 

377 self.tap_schema_setup = TapSchemaSqliteSetup(TEST_COMPOSITE_KEYS, context={"id_generation": True}) 

378 

379 # Fetch the keys and key_columns data from the TAP_SCHEMA tables. 

380 keys_table = self.tap_schema_setup.mgr["keys"] 

381 key_columns_table = self.tap_schema_setup.mgr["key_columns"] 

382 with self.tap_schema_setup.engine.connect() as connection: 

383 key_columns_result = connection.execute(select(key_columns_table)) 

384 self.key_columns_data = [row._asdict() for row in key_columns_result] 

385 

386 keys_result = connection.execute(select(keys_table)) 

387 self.keys_data = [row._asdict() for row in keys_result] 

388 

389 def test_keys(self) -> None: 

390 """Test that composite keys are handled correctly by inspecting the 

391 data in the generated TAP_SCHEMA ``keys`` table. 

392 """ 

393 print(f"\nComposite keys data: {self.keys_data}") 

394 

395 self.assertEqual(len(self.keys_data), 1) 

396 

397 self.assertEqual( 

398 self.keys_data[0], 

399 { 

400 "key_id": "fk_composite", 

401 "from_table": "test_composite_keys.table1", 

402 "target_table": "test_composite_keys.table2", 

403 "utype": "ForeignKey", 

404 "description": "Composite foreign key from table1 to table2", 

405 }, 

406 ) 

407 

408 def test_key_columns(self) -> None: 

409 """Test that composite keys are handled correctly by inspecting the 

410 data in the generated TAP_SCHEMA ``key_columns`` table. 

411 """ 

412 print(f"\nComposite key columns data: {self.key_columns_data}") 

413 

414 self.assertEqual(len(self.key_columns_data), 2) 

415 

416 key_columns_row1 = self.key_columns_data[0] 

417 self.assertEqual( 

418 key_columns_row1, {"key_id": "fk_composite", "from_column": "id1", "target_column": "id1"} 

419 ) 

420 

421 key_columns_row2 = self.key_columns_data[1] 

422 self.assertEqual( 

423 key_columns_row2, {"key_id": "fk_composite", "from_column": "id2", "target_column": "id2"} 

424 ) 

425 

426 

427if __name__ == "__main__": 427 ↛ 428line 427 didn't jump to line 428 because the condition on line 427 was never true

428 unittest.main()