File size: 5,524 Bytes
92ea1b5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
import os
import sys
import tempfile
import unittest
from unittest.mock import MagicMock, patch

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))


class FakeModel:
    "Minimal model stub for testing without GPU"
    def __init__(self):
        self.alphabet = {c: i + 3 for i, c in enumerate("ACDEFGHIKLMNPQRSTVWY")}
        self.alphabet.update({'<pad>': 1, '<eos>': 2, '<cls>': 0})


class TestAppValidation(unittest.TestCase):
    "app() input guards raise gradio.Error on bad input"

    @patch('data.ModelFactory', return_value=FakeModel())
    def test_empty_sequence_raises(self, _mock_factory):
        from app import app
        from gradio import Error
        with self.assertRaises(Error) as ctx:
            app("", "scan", "test/model", True)
        self.assertIn("empty", str(ctx.exception).lower())

    @patch('data.ModelFactory', return_value=FakeModel())
    def test_empty_substitutions_raises(self, _mock_factory):
        from app import app
        from gradio import Error
        with self.assertRaises(Error) as ctx:
            app("MVEQYLLEAI", "", "test/model", True)
        self.assertIn("substitution", str(ctx.exception).lower())


class TestDmsLayout(unittest.TestCase):
    "ncols/nrows layout calculation logic from _render_dms"

    def _calc_layout(self, num_cols):
        """Replicate the ncols selection logic from data.py"""
        from math import ceil
        ncols = min([d for d in range(1, num_cols + 1) if num_cols % d == 0 and 30 <= d <= 60] or [60],
                    key=lambda x: abs(x - 60))
        nrows = ceil(num_cols / ncols)
        while num_cols / ncols < nrows and ncols > 45 and ncols * nrows >= num_cols:
            ncols -= 1
        ncols += 1
        return ncols, nrows

    def test_short_sequence(self):
        # 7-col sequence (like examples in app)
        ncols, nrows = self._calc_layout(7)
        self.assertEqual(nrows, 1)
        self.assertGreater(ncols, 0)

    def test_medium_sequence(self):
        ncols, nrows = self._calc_layout(100)
        self.assertGreaterEqual(ncols, 30)
        self.assertTrue(ncols * nrows >= 100)  # grid must cover all columns

    def test_long_sequence_fits_grid(self):
        ncols, nrows = self._calc_layout(300)
        self.assertTrue(ncols * nrows >= 300)

    def test_very_long_sequence(self):
        ncols, nrows = self._calc_layout(600)
        self.assertTrue(ncols * nrows >= 600)


class TestDmsRender(unittest.TestCase):
    "Full _render_dms pipeline produces PNG file"

    @patch('data.ModelFactory', return_value=FakeModel())
    def test_render_creates_png(self, _mock_factory):
        from data import Data
        import pandas as pd
        with tempfile.TemporaryDirectory() as tmpdir:
            out_path = os.path.join(tmpdir, 'test_heatmap.png')
            csv_path = os.path.join(tmpdir, 'test_out.csv')
            d = object.__new__(Data)
            d.model_name = 'test'
            d.seq = "MVEQYLL"
            d.mode = 'DMS'
            d.resi = list(range(1, len(d.seq) + 1))
            # Build a minimal scored output matching DMS expectations
            AA = "ACDEFGHIKLMNPQRSTVWY"
            rows = []
            for i, src_c in enumerate(d.seq, 1):
                for trg_c in AA.replace(src_c, ''):
                    rows.append(f"{src_c}{i}{trg_c}")
            scores = [float(i % 5 - 2) for i in range(len(rows))]
            d.out = pd.DataFrame({'0': rows, 'test': scores})
            d.out_img_path = out_path
            d.out_csv = csv_path
            d._render_dms()
            self.assertTrue(os.path.exists(out_path), f"PNG not created at {out_path}")
            self.assertTrue(os.path.getsize(out_path) > 0, "PNG is empty")


class TestParseOutputDispatch(unittest.TestCase):
    "parse_output routes to correct handler per mode"

    @patch('data.ModelFactory', return_value=FakeModel())
    def test_dms_dispatch_calls_render(self, _mock_factory):
        from data import Data
        import pandas as pd
        with tempfile.TemporaryDirectory() as tmpdir:
            d = object.__new__(Data)
            d.model_name = 'test'
            d.mode = 'DMS'
            d.seq = "MV"
            d.resi = [1, 2]
            AA = "ACDEFGHIKLMNPQRSTVWY"
            rows = []
            for src_c in "MV":
                pos = {"M": 1, "V": 2}[src_c]
                for trg_c in AA.replace(src_c, ''):
                    rows.append(f"{src_c}{pos}{trg_c}")
            d.out = pd.DataFrame({'0': rows, 'test': list(range(len(rows)))})
            d.out_img_path = os.path.join(tmpdir, 't.png')
            d.out_csv = os.path.join(tmpdir, 't.csv')
            d.parse_output()
            self.assertTrue(os.path.exists(d.out_img_path))

    @patch('data.ModelFactory', return_value=FakeModel())
    def test_mut_dispatch_sorts_and_styles(self, _mock_factory):
        from data import Data
        import pandas as pd
        with tempfile.TemporaryDirectory() as tmpdir:
            d = object.__new__(Data)
            d.model_name = 'test'
            d.mode = 'MUT'
            d.out = pd.DataFrame({'0': ['V2A', 'E5K'], 'test': [3.0, -1.0]})
            d.out_csv = os.path.join(tmpdir, 'm.csv')
            d.out_img_path = os.path.join(tmpdir, 'm.png')
            d.parse_output()
            # MUT sorts descending by score; V2A (3.0) should be first
            self.assertEqual(d.out.iloc[0]['0'], 'V2A')
            self.assertIsNotNone(d.out_table)


if __name__ == '__main__':
    unittest.main()