test/test_app.py

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()