Coverage for tests/test_bedfile.py: 87%

1import os 

2 

3import numpy as np 

4import pytest 

5from gtars.models import RegionSet as GRegionSet 

6from qdrant_client.models import PointIdsList 

7from sqlalchemy.orm import Session 

8from sqlalchemy.sql import select 

9 

10from bbconf.bbagent import BedBaseAgent 

11from bbconf.const import DEFAULT_LICENSE 

12from bbconf.db_utils import Bed, Files 

13from bbconf.exceptions import BedFIleExistsError, BEDFileNotFoundError 

14from bbconf.modules.bedfiles import _bed_id_from_point_id 

15 

16from .conftest import DATA_PATH, SERVICE_UNAVAILABLE, get_bbagent 

17from .utils import BED_TEST_ID, BEDSET_TEST_ID, ContextManagerDBTesting 

18 

19 

20@pytest.mark.skipif(SERVICE_UNAVAILABLE, reason="Database is not available") 

21def test_bb_database(): 

22 agent = get_bbagent() 

23 assert isinstance(agent, BedBaseAgent) 

24 

25 

26@pytest.mark.skipif(SERVICE_UNAVAILABLE, reason="Database is not available") 

27class Test_BedFile_Agent: 

28 def test_upload(self, bbagent_obj, example_dict, mocker): 

29 upload_s3_mock = mocker.patch( 

30 "bbconf.config_parser.bedbaseconfig.BedBaseConfig.upload_s3", 

31 return_value=True, 

32 ) 

33 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=False): 

34 bbagent_obj.bed.add(**example_dict) 

35 

36 assert upload_s3_mock.called 

37 assert bbagent_obj.bed.exists(example_dict["identifier"]) 

38 

39 def test_upload_exists(self, bbagent_obj, example_dict, mocker): 

40 mocker.patch( 

41 "bbconf.config_parser.bedbaseconfig.BedBaseConfig.upload_s3", 

42 return_value=True, 

43 ) 

44 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=False): 

45 bbagent_obj.bed.add(**example_dict) 

46 with pytest.raises(BedFIleExistsError): 

47 bbagent_obj.bed.add(**example_dict) 

48 

49 def test_add_nofail(self, bbagent_obj, example_dict, mocker): 

50 mocker.patch( 

51 "bbconf.config_parser.bedbaseconfig.BedBaseConfig.upload_s3", 

52 return_value=True, 

53 ) 

54 

55 example_dict["nofail"] = True 

56 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=False): 

57 bbagent_obj.bed.add(**example_dict) 

58 bbagent_obj.bed.add(**example_dict) 

59 assert bbagent_obj.bed.exists(example_dict["identifier"]) 

60 

61 def test_get_all(self, bbagent_obj): 

62 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

63 return_result = bbagent_obj.bed.get(BED_TEST_ID, full=True) 

64 assert return_result is not None 

65 assert return_result.files is not None 

66 assert return_result.plots is not None 

67 

68 assert return_result.genome_alias == "hg38" 

69 assert return_result.stats.number_of_regions == 1 

70 

71 assert return_result.files.bed_file is not None 

72 assert return_result.plots.chrombins is not None 

73 assert return_result.license_id == DEFAULT_LICENSE 

74 

75 def test_get_all_bedsets_bedfile_count(self, bbagent_obj): 

76 with ContextManagerDBTesting( 

77 config=bbagent_obj.config, add_data=True, bedset=True 

78 ): 

79 return_result = bbagent_obj.bed.get(BED_TEST_ID, full=True) 

80 

81 assert len(return_result.bedsets) == 1 

82 assert return_result.bedsets[0].id == BEDSET_TEST_ID 

83 assert return_result.bedsets[0].bedfile_count == 1 

84 

85 def test_get_all_not_found(self, bbagent_obj): 

86 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

87 return_result = bbagent_obj.bed.get(BED_TEST_ID, full=False) 

88 

89 assert return_result is not None 

90 assert return_result.files is None 

91 assert return_result.plots is None 

92 assert return_result.stats is None 

93 

94 assert return_result.genome_alias == "hg38" 

95 assert return_result.id == BED_TEST_ID 

96 

97 def test_get_stats(self, bbagent_obj): 

98 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

99 return_result = bbagent_obj.bed.get_stats(BED_TEST_ID) 

100 

101 assert return_result is not None 

102 assert return_result.number_of_regions == 1 

103 

104 def test_get_files(self, bbagent_obj): 

105 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

106 return_result = bbagent_obj.bed.get_files(BED_TEST_ID) 

107 

108 assert return_result is not None 

109 assert return_result.bed_file.path is not None 

110 

111 def test_get_plots(self, bbagent_obj): 

112 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

113 return_result = bbagent_obj.bed.get_plots(BED_TEST_ID) 

114 

115 assert return_result is not None 

116 assert return_result.chrombins is not None 

117 

118 def test_get_objects(self, bbagent_obj): 

119 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

120 return_result = bbagent_obj.bed.get_objects(BED_TEST_ID) 

121 

122 assert "bed_file" in return_result 

123 assert "chrombins" in return_result 

124 

125 def test_get_classification(self, bbagent_obj): 

126 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

127 return_result = bbagent_obj.bed.get_classification(BED_TEST_ID) 

128 

129 assert return_result is not None 

130 assert return_result.bed_compliance == "bed6+4" 

131 

132 def test_get_list(self, bbagent_obj): 

133 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

134 return_result = bbagent_obj.bed.get_ids_list(limit=100, offset=0) 

135 

136 assert len(return_result.results) == 1 

137 assert return_result.count == 1 

138 assert return_result.results[0].id == BED_TEST_ID 

139 assert return_result.limit == 100 

140 assert return_result.offset == 0 

141 

142 def test_get_list_genome_true(self, bbagent_obj): 

143 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

144 return_result = bbagent_obj.bed.get_ids_list( 

145 limit=100, offset=0, genome="hg38" 

146 ) 

147 

148 assert len(return_result.results) == 1 

149 assert return_result.count == 1 

150 assert return_result.results[0].id == BED_TEST_ID 

151 assert return_result.limit == 100 

152 assert return_result.offset == 0 

153 

154 def test_get_list_genome_false(self, bbagent_obj): 

155 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

156 return_result = bbagent_obj.bed.get_ids_list( 

157 limit=100, offset=0, genome="hg381" 

158 ) 

159 

160 assert len(return_result.results) == 0 

161 assert return_result.count == 0 

162 

163 def test_get_list_bed_compliance_true(self, bbagent_obj): 

164 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

165 return_result = bbagent_obj.bed.get_ids_list( 

166 limit=100, offset=0, bed_compliance="bed6+4" 

167 ) 

168 

169 assert len(return_result.results) == 1 

170 assert return_result.count == 1 

171 assert return_result.results[0].id == BED_TEST_ID 

172 assert return_result.limit == 100 

173 assert return_result.offset == 0 

174 

175 def test_get_list_bed_compliance_false(self, bbagent_obj): 

176 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

177 return_result = bbagent_obj.bed.get_ids_list( 

178 limit=100, offset=0, bed_compliance="bed6+5" 

179 ) 

180 

181 assert len(return_result.results) == 0 

182 assert return_result.count == 0 

183 

184 def test_get_list_bed_offset(self, bbagent_obj): 

185 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

186 return_result = bbagent_obj.bed.get_ids_list( 

187 limit=100, 

188 offset=1, 

189 ) 

190 

191 assert len(return_result.results) == 0 

192 # count is only computed on the first page 

193 assert return_result.count is None 

194 assert return_result.offset == 1 

195 

196 def test_get_list_after_walks_every_id_once(self, bbagent_obj): 

197 extra_ids = [f"{i:032x}" for i in range(1, 8)] 

198 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

199 with Session(bbagent_obj.config.db_engine.engine) as session: 

200 for bed_id in extra_ids: 

201 session.add(Bed(id=bed_id, genome_alias="hg38", processed=False)) 

202 session.commit() 

203 

204 expected = sorted(extra_ids + [BED_TEST_ID]) 

205 

206 seen = [] 

207 after = None 

208 while True: 

209 page = bbagent_obj.bed.get_ids_list(limit=3, after=after) 

210 if after is None: 

211 assert page.count == len(expected) 

212 else: 

213 assert page.count is None 

214 assert page.offset == 0 

215 if not page.results: 

216 break 

217 seen.extend(r.id for r in page.results) 

218 after = page.results[-1].id 

219 

220 assert seen == expected 

221 assert len(seen) == len(set(seen)) 

222 

223 def test_get_list_after_ignores_offset(self, bbagent_obj): 

224 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

225 with Session(bbagent_obj.config.db_engine.engine) as session: 

226 session.add(Bed(id="0" * 32, genome_alias="hg38", processed=False)) 

227 session.commit() 

228 

229 return_result = bbagent_obj.bed.get_ids_list( 

230 limit=100, offset=5, after="0" * 32 

231 ) 

232 

233 assert [r.id for r in return_result.results] == [BED_TEST_ID] 

234 assert return_result.count is None 

235 assert return_result.offset == 0 

236 

237 def test_bed_delete(self, bbagent_obj, mocker): 

238 mocker.patch("bbconf.config_parser.bedbaseconfig.BedBaseConfig.delete_s3") 

239 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

240 bbagent_obj.bed.delete(BED_TEST_ID) 

241 

242 assert not bbagent_obj.bed.exists(BED_TEST_ID) 

243 

244 with Session(bbagent_obj.config.db_engine.engine) as session: 

245 result = session.scalar(select(Bed).where(Bed.id == BED_TEST_ID)) 

246 assert result is None 

247 

248 result = session.scalars(select(Files)) 

249 assert len([k for k in result]) == 0 

250 

251 def test_bed_delete_not_found(self, bbagent_obj): 

252 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

253 with pytest.raises(BEDFileNotFoundError): 

254 bbagent_obj.bed.delete("not_found") 

255 

256 def test_bed_update(self, bbagent_obj): 

257 

258 # TODO: has to be expanded 

259 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

260 bed_file = bbagent_obj.bed.get(BED_TEST_ID, full=True) 

261 # assert bed_file.annotation.model_dump(exclude_defaults=True) == {} 

262 assert bed_file.annotation.cell_line is None 

263 

264 new_metadata = { 

265 "cell_line": "K562", 

266 "tissue": "blood", 

267 } 

268 bbagent_obj.bed.update( 

269 identifier=BED_TEST_ID, 

270 metadata=new_metadata, 

271 upload_qdrant=False, 

272 upload_s3=False, 

273 ) 

274 

275 new_bed_file = bbagent_obj.bed.get(BED_TEST_ID, full=True) 

276 

277 assert new_bed_file.annotation.cell_line == "K562" 

278 

279 def test_get_unprocessed(self, bbagent_obj): 

280 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

281 return_result = bbagent_obj.bed.get_unprocessed(limit=100, offset=0) 

282 

283 assert return_result.count == 1 

284 assert return_result.results[0].id == BED_TEST_ID 

285 

286 def test_get_missing_plots(self, bbagent_obj): 

287 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

288 return_result = bbagent_obj.bed.get_missing_plots( 

289 "tss_distance", limit=100, offset=0 

290 ) 

291 

292 assert return_result[0] == BED_TEST_ID 

293 

294 

295def test_bed_id_from_point_id(): 

296 dashed = "-".join( 

297 [ 

298 BED_TEST_ID[:8], 

299 BED_TEST_ID[8:12], 

300 BED_TEST_ID[12:16], 

301 BED_TEST_ID[16:20], 

302 BED_TEST_ID[20:], 

303 ] 

304 ) 

305 assert _bed_id_from_point_id(dashed) == BED_TEST_ID 

306 assert _bed_id_from_point_id(BED_TEST_ID) == BED_TEST_ID 

307 

308 

309def _ml_unavailable() -> bool: 

310 """True unless ML models loaded and the Qdrant server answers.""" 

311 if SERVICE_UNAVAILABLE: 

312 return True 

313 config = get_bbagent().config 

314 if getattr(config, "bivec_search_interface", None) is None: 

315 return True 

316 try: 

317 config.qdrant_client.get_collections() 

318 except Exception: 

319 return True 

320 return False 

321 

322 

323ML_UNAVAILABLE = _ml_unavailable() 

324QDRANT_TEXT_POINT_ID = "00000000-0000-0000-0000-000000000001" 

325 

326 

327@pytest.fixture() 

328def qdrant_bed_point(bbagent_obj): 

329 """Load one BED vector and one matching text vector into Qdrant. 

330 

331 The text vector is the real dense embedding of a metadata phrase, and its 

332 payload points at BED_TEST_ID, so a bivec text search can find the BED. 

333 """ 

334 config = bbagent_obj.config 

335 qdrant_cfg = config.config.qdrant 

336 file_dim = config.qdrant_client.get_collection( 

337 qdrant_cfg.file_collection 

338 ).config.params.vectors.size 

339 config.qdrant_file_backend.load( 

340 ids=[BED_TEST_ID], 

341 vectors=np.random.default_rng(0).random((1, file_dim)), 

342 payloads=[{"description": "test bed"}], 

343 ) 

344 text_vec = next(iter(config.dense_encoder.embed(["CTCF ChIP-seq in K562"]))) 

345 config._qdrant_text_backend.load( 

346 ids=[QDRANT_TEXT_POINT_ID], 

347 vectors=np.array([text_vec]), 

348 payloads=[{"matched_files": [BED_TEST_ID]}], 

349 ) 

350 yield 

351 config.qdrant_client.delete( 

352 collection_name=qdrant_cfg.file_collection, 

353 points_selector=PointIdsList(points=[BED_TEST_ID]), 

354 ) 

355 config.qdrant_client.delete( 

356 collection_name=qdrant_cfg.text_collection, 

357 points_selector=PointIdsList(points=[QDRANT_TEXT_POINT_ID]), 

358 ) 

359 

360 

361def _file_points_count(bbagent_obj) -> int: 

362 return bbagent_obj.config.qdrant_client.count( 

363 collection_name=bbagent_obj.config.config.qdrant.file_collection, 

364 exact=True, 

365 ).count 

366 

367 

368@pytest.mark.skipif(ML_UNAVAILABLE, reason="ML models or Qdrant server unavailable") 

369class TestVectorSearch: 

370 def test_text_to_bed_search(self, bbagent_obj, qdrant_bed_point): 

371 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

372 return_result = bbagent_obj.bed.text_to_bed_search("CTCF ChIP-seq") 

373 

374 assert return_result.count == 1 

375 assert len(return_result.results) == 1 

376 # Qdrant returns dashed UUIDs; search results use the plain BED id. 

377 assert return_result.results[0].id == BED_TEST_ID 

378 assert return_result.results[0].metadata.id == BED_TEST_ID 

379 

380 def test_bed_to_bed_search(self, bbagent_obj, qdrant_bed_point): 

381 if bbagent_obj.config.b2b_search_interface is None: 

382 pytest.skip("Region2vec model unavailable") 

383 region_set = GRegionSet( 

384 os.path.join(DATA_PATH, "files", f"{BED_TEST_ID}.bed.gz") 

385 ) 

386 with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): 

387 return_result = bbagent_obj.bed.bed_to_bed_search(region_set, limit=5) 

388 

389 assert len(return_result.results) == 1 

390 assert return_result.results[0].id == BED_TEST_ID 

391 assert return_result.results[0].metadata.id == BED_TEST_ID 

392 

393 def test_delete_qdrant_point(self, bbagent_obj, qdrant_bed_point): 

394 assert _file_points_count(bbagent_obj) == 1 

395 bbagent_obj.bed.delete_qdrant_point(BED_TEST_ID) 

396 assert _file_points_count(bbagent_obj) == 0 

397 

398 def test_qdrant_collections_created(self, bbagent_obj): 

399 """Agent init creates every configured Qdrant collection.""" 

400 qdrant_cfg = bbagent_obj.config.config.qdrant 

401 for collection in ( 

402 qdrant_cfg.file_collection, 

403 qdrant_cfg.text_collection, 

404 qdrant_cfg.hybrid_collection, 

405 ): 

406 assert bbagent_obj.config.qdrant_client.collection_exists(collection)