| import pytest |
| from pathlib import Path |
| from model.models.results import PredictionResult |
| from model.prediction_result_file import PredictionResultsFile |
| from model import errors |
|
|
| @pytest.fixture |
| def create_n_mock_results(): |
| """ |
| Returns a fixture that creates n mock results |
| """ |
| def _create_n_mock_results(n): |
| return [ |
| PredictionResult( |
| pdb_path=Path(f"mock_path_{i}.pdb"), |
| chain_id=f"mock_chain_id_{i}", |
| sequence_md5=f"mock_md5_{i}", |
| ndom=1, |
| nres=1234, |
| chopping="12-34_56-78,90-123", |
| confidence=0.0123, |
| time_sec=0.1234, |
| ) for i in range(n) |
| ] |
| return _create_n_mock_results |
|
|
|
|
| def count_lines(path): |
| with open(path) as f: |
| return sum(1 for _ in f) |
|
|
| def test_chunked_results(tmp_path, create_n_mock_results): |
|
|
| expected_path = tmp_path / "expected_results.tsv" |
| result_file = PredictionResultsFile(expected_path, chunk_size=20) |
| results = create_n_mock_results(50) |
| expected_headers = '\t'.join([ |
| 'chain_id', 'sequence_md5', 'nres', 'ndom', 'chopping', 'confidence', 'time_sec' |
| ]) |
|
|
| assert expected_path.exists() == False |
|
|
| results_before_chunk = results[0:19] |
| assert len(results_before_chunk) == 19 |
| result_file.add_results(results_before_chunk) |
| assert expected_path.exists() == False |
|
|
| results_after_chunk = results[19:20] |
| assert len(results_after_chunk) == 1 |
| result_file.add_results(results_after_chunk) |
| assert expected_path.exists() == True |
|
|
| lines = expected_path.read_text().split('\n') |
| firstline = lines[0] |
| assert firstline == expected_headers |
| assert 'mock_chain_id_0' in lines[1] |
| assert 'mock_chain_id_19' in lines[20] |
| assert count_lines(expected_path) == 21 |
| |
|
|
| def test_check_add_repeated_result_raises_error(tmp_path, create_n_mock_results): |
|
|
| expected_path = tmp_path / "expected_results.tsv" |
| result_file = PredictionResultsFile(expected_path, chunk_size=20) |
| results = create_n_mock_results(10) |
|
|
| assert expected_path.exists() == False |
|
|
| result_file.add_result(results[0]) |
| with pytest.raises(errors.PredictionResultExistsError) as err: |
| result_file.add_result(results[0]) |
|
|
|
|
|
|
| @pytest.mark.parametrize("allow_append,exception", [ |
| (True, None), |
| (False, errors.FileExistsError), |
| ]) |
| def test_check_appending_to_existing_file(tmp_path, create_n_mock_results, allow_append, exception): |
|
|
| expected_path = tmp_path / "expected_results.tsv" |
| results = create_n_mock_results(60) |
|
|
| def get_result_file(): |
| return PredictionResultsFile(expected_path, chunk_size=20, allow_append=allow_append) |
|
|
| f1 = get_result_file() |
|
|
| assert expected_path.exists() == False |
| f1.add_results(results[0:25]) |
| assert expected_path.exists() == True |
| assert count_lines(expected_path) == 21 |
| f1.flush() |
| assert count_lines(expected_path) == 26 |
|
|
| def append_results(): |
| f2 = get_result_file() |
| f2.add_results(results[25:50]) |
| f2.flush() |
|
|
| if exception is None: |
| append_results() |
| count_lines(expected_path) == 51 |
| f3 = get_result_file() |
| assert f3.has_result(results[0]) is True |
| assert f3.has_result(results[59]) is False |
| else: |
| with pytest.raises(exception) as err: |
| append_results() |
|
|
