| from pathlib import Path |
| import io |
| from model.models.results import PredictionResult |
| from scripts.get_predictions import get_csv_writer, write_csv_results |
|
|
|
|
| def test_write_csv_results(): |
|
|
| mock_pdb_path = Path("mock.pdb") |
| chain_id = 'mock_id' |
| expected_sequence_md5 = 'mock_md5' |
| expected_chopping_str = '189-237,238-353,381-505,524-658,687-826' |
| expected_ndom = len(expected_chopping_str.split(',')) |
| expected_nres = 1234 |
| expected_confidence = 0.8910 |
| expected_time_sec = 0.1234 |
|
|
| result = PredictionResult( |
| pdb_path=mock_pdb_path, |
| chain_id=chain_id, |
| sequence_md5=expected_sequence_md5, |
| ndom=expected_ndom, |
| nres=expected_nres, |
| chopping=expected_chopping_str, |
| confidence=expected_confidence, |
| time_sec=expected_time_sec, |
| ) |
|
|
| fp = io.StringIO() |
| csv_writer = get_csv_writer(fp) |
| write_csv_results(csv_writer, [result]) |
| fp.flush() |
| fp.seek(0) |
|
|
| expected_cols = [chain_id, expected_sequence_md5, expected_nres, expected_ndom, |
| expected_chopping_str, expected_confidence, expected_time_sec] |
| actual_text = fp.read().replace('\r\n', '\n') |
| assert actual_text == '\t'.join([str(col) for col in expected_cols]) + '\n' |
|
|
