add fasta support; add unittest suite

.gitignore

@@ -2,3 +2,4 @@ Dockerfile
 *.ipynb
 out.*
 */
+__pycache__/

data.py

@@ -15,8 +15,13 @@ class Data:
     AA = "ACDEFGHIKLMNPQRSTVWY"
 
     def parse_seq(self, src: str):
-        "Parse input sequence"
-        self.seq = src.strip().upper().replace('\n', '')
+        "Parse input sequence (plain or FASTA)"
+        lines = src.strip().splitlines()
+        if lines and lines[0].startswith('>'):
+            lines = lines[1:]
+        self.seq = ''.join(lines).upper().replace('\n', '').replace(' ', '')
+        if not self.seq:
+            raise RuntimeError("Sequence is empty")
         if not all(x in self.model.alphabet for x in self.seq):
             raise RuntimeError(f"Unsupported characters in sequence: {''.join(x for x in self.seq if x not in self.model.alphabet)}")
 
@@ -35,7 +40,7 @@ class Data:
                     self.sub.append(f"{src_c}{resi}{trg_c}")
                     self.resi.append(resi)
         elif all(match(r'\d+', x) for x in self.trg):
-            self.mode = 'DMS'
+            self.mode = 'SMS'
             trh = []
             for resi in map(int, self.trg):
                 if resi < 1 or resi > len(self.seq):
@@ -59,7 +64,7 @@ class Data:
                     raise RuntimeError(f"Unrecognised input substitution: {self.seq[int(''.join(resi_str))]}{int(''.join(resi_str))} /= {s}{int(''.join(resi_str))}")
             self.trg = trh
         else:
-            self.mode = 'TMS'
+            self.mode = 'DMS'
             self.trg = []
             for resi, src_c in enumerate(self.seq, 1):
                 for trg_c in self.AA.replace(src_c, ''):
@@ -84,11 +89,11 @@ class Data:
 
     def parse_output(self) -> None:
         "Format output data for visualisation"
-        if self.mode == "TMS":
-            self._render_tms()
+        if self.mode == "DMS":
+            self._render_dms()
             self.out.to_csv(self.out_csv, float_format='%.2f')
-        elif self.mode == "DMS":
-            self._sort_dms()
+        elif self.mode == "SMS":
+            self._sort_sms()
             self._style_and_save()
         elif self.mode == "MUT":
             self.out = self.out.sort_values(self.model_name, ascending=False)
@@ -103,16 +108,16 @@ class Data:
                 .background_gradient(cmap="RdYlGn", vmax=8, vmin=-8).hide(axis=0).hide(axis=1))
         self.out.to_csv(self.out_csv, float_format='%.2f', index=False, header=False)
 
-    def _sort_dms(self):
-        "Sort DMS output by residue then score, top 19 per position, reshaped to columns"
+    def _sort_sms(self):
+        "Sort SMS output by residue then score, top 19 per position, reshaped to columns"
         self.out = (self.out.assign(resi=self.out['0'].str.extract(r'(\d+)', expand=False).astype(int))
                         .sort_values(["resi", self.model_name], ascending=[True,False])
                         .groupby(["resi"]).head(19).drop(["resi"], axis=1))
         self.out = pd.concat([self.out.iloc[19*x:19*(x+1)].reset_index(drop=True) for x in range(self.out.shape[0]//19)]
                              , axis=1).set_axis(range(self.out.shape[0]//19*2), axis="columns")
 
-    def _render_tms(self):
-        "Build TMS heatmap from scored mutations"
+    def _render_dms(self):
+        "Build DMS heatmap from scored mutations"
         # Group by residue, keep top 19 (all alternatives)
         grouped = (self.out.assign(resi=self.resi_cycle())
                          .groupby("resi").head(19))
@@ -147,7 +152,7 @@ class Data:
         return (self.resi * (len(self.out)//len(self.resi) + 1))[:len(self.out)]
 
     def _plot_heatmap(self, ncols, nrows):
-        "Render TMS heatmap to PNG"
+        "Render DMS heatmap to PNG"
         kw = dict(cmap="RdBu", cbar=False, square=True, xticklabels=1, yticklabels=1
                   , center=0, fmt='s', annot_kws={"size": "xx-large"})
         annotate = lambda df: df.map(lambda x: ' ' if x != 0 else '\u00b7')
@@ -191,7 +196,7 @@ class Data:
 
     @property
     def image(self):
-        "return PNG path (TMS) or Styler object (DMS/MUT)"
+        "return PNG path (DMS) or Styler object (SMS/MUT)"
         if self.out_table is not None:
             return self.out_table
         return self.out_img_path

instructions.md

@@ -6,7 +6,7 @@ If the server remains idle for a period, it will enter standby mode. Running a c
 
 ## Input
 
-**Sequence**: Enter the full amino acid sequence to be analyzed in the **Sequence** text box.  
+**Sequence**: Enter the full amino acid sequence to be analyzed in the **Sequence** text box. Plain sequences or FASTA format (with `>` header line) are accepted.  
 Note: While jolly characters (e.g., `-X.B`) can be included, they currently cannot be visualized.
 
 **Substitutions**: Specify the substitutions you wish to test in the **Substitutions** box. The tool supports three running modes based on your input:

model.py

@@ -64,8 +64,8 @@ class ESMModel:
 
         # Vectorized scoring: extract WT chars, indices, mutant chars from mutation strings
         muts = data.sub['0'].values
-        wt_ids = torch.tensor([self[c] for c in muts], device=self.device)
-        idxs = torch.tensor([[int(d) for d in m if d.isdigit()] for m in muts], device=self.device).squeeze(-1) - 1
+        wt_ids = torch.tensor([self[m[0]] for m in muts], device=self.device)
+        idxs = torch.tensor([int(''.join(c for c in m if c.isdigit())) for m in muts], device=self.device) - 1
         mt_ids = torch.tensor([self[m[-1]] for m in muts], device=self.device)
         scores = token_probs.to(self.device)[0, 1 + idxs, mt_ids] - token_probs.to(self.device)[0, 1 + idxs, wt_ids]
         data.out[self.model_name] = scores.cpu().numpy()

test/test_app.py

@@ -0,0 +1,144 @@
+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()

test/test_data.py

@@ -0,0 +1,297 @@
+import os
+import sys
+import tempfile
+import unittest
+from unittest.mock import MagicMock, patch
+
+# Ensure project root is on path
+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 TestParseSeq(unittest.TestCase):
+    "Sequence parsing (plain and FASTA)"
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_plain_sequence(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.parse_seq("MVEQYLLEAI")
+        self.assertEqual(d.seq, "MVEQYLLEAI")
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_fasta_single_line(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.parse_seq(">my protein\nMVEQYLLEAI")
+        self.assertEqual(d.seq, "MVEQYLLEAI")
+        self.assertFalse(d.seq.startswith('>'))
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_fasta_multi_line(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.parse_seq(">seq1\nMV EQ YL LE AI")
+        self.assertEqual(d.seq, "MVEQYLLEAI")
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_lowercase_converted_to_upper(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.parse_seq("mveqylleai")
+        self.assertEqual(d.seq, "MVEQYLLEAI")
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_whitespace_stripped(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.parse_seq("  MV EQ YL \n LE AI  ")
+        self.assertEqual(d.seq, "MVEQYLLEAI")
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_invalid_characters_raises(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        with self.assertRaises(RuntimeError):
+            d.parse_seq("MVXYZ")
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_empty_sequence_raises(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        with self.assertRaises(RuntimeError):
+            d.parse_seq("")
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_whitespace_only_raises(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        with self.assertRaises(RuntimeError):
+            d.parse_seq("   \n  ")
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_fasta_header_only_raises(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        with self.assertRaises(RuntimeError):
+            d.parse_seq(">just a header")
+
+
+class TestParseSub(unittest.TestCase):
+    "Substitution parsing and mode detection"
+
+    def setUp(self):
+        self.seq = "MVEQYLL"
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_dms_mode(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.seq = self.seq
+        d.parse_sub("2 5")
+        self.assertEqual(d.mode, 'SMS')
+        self.assertEqual(len(d.resi), 2)
+        self.assertIn(2, d.resi)
+        self.assertIn(5, d.resi)
+        # Each position has 19 alternatives (20 AA minus WT)
+        self.assertEqual(len(d.sub), 38)
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_mut_explicit_mode(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.seq = self.seq
+        d.parse_sub("V2A E3K")
+        self.assertEqual(d.mode, 'MUT')
+        self.assertEqual(len(d.sub), 2)
+        self.assertEqual(list(d.sub['0']), ['V2A', 'E3K'])
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_mut_seq_vs_seq_mode(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.seq = self.seq
+        d.parse_sub("MVEQYAL")  # same length, differs at pos 6
+        self.assertEqual(d.mode, 'MUT')
+        self.assertTrue(any('L6A' in str(s) for s in d.sub['0']))
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_tms_fallback_mode(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.seq = self.seq
+        d.parse_sub("deep mutational scanning")
+        self.assertEqual(d.mode, 'DMS')
+        # All positions x all alternatives
+        self.assertEqual(len(d.resi), len(self.seq))
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_dms_position_out_of_range(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.seq = self.seq
+        with self.assertRaises(RuntimeError):
+            d.parse_sub("999")
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_mut_wrong_wt_raises(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.seq = self.seq
+        # V2A but position 2 is V — correct; try A2K where pos 2 is V not A
+        with self.assertRaises(RuntimeError):
+            d.parse_sub("A2K")
+
+
+class TestSmsSort(unittest.TestCase):
+    "SMS output sorting and reshaping"
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_sort_preserves_top_19_per_position(self, mock_factory):
+        from data import Data
+        import pandas as pd
+        d = object.__new__(Data)
+        d.model_name = 'test'
+        # _sort_sms requires exactly 19 rows per position; build minimal valid input
+        AA = "ACDEFGHIKLMNPQRSTVWY"
+        subs_pos2 = [f'V2{a}' for a in AA.replace('V', '')]   # 19 alternatives (skip WT V)
+        subs_pos5 = [f'L5{a}' for a in AA.replace('L', '')]   # 19 alternatives (skip WT L)
+        scores_p2 = list(range(19))
+        scores_p5 = list(range(18, -1, -1))
+        d.out = pd.DataFrame({
+            '0': subs_pos2 + subs_pos5,
+            'test': scores_p2[:19] + scores_p5
+        })
+        d.resi = [2, 5]
+        d._sort_sms()
+        self.assertEqual(d.out.shape[0], 19)  # 19 rows after reshape
+
+
+class TestResidueCycle(unittest.TestCase):
+    "resi_cycle helper"
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_resi_cycles_correctly(self, _mock_factory):
+        from data import Data
+        import pandas as pd
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.resi = [1, 3, 5]
+        d.out = pd.DataFrame({'x': range(9)})
+        cycle = d.resi_cycle()
+        self.assertEqual(cycle, [1, 3, 5, 1, 3, 5, 1, 3, 5])
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_resi_truncates_to_length(self, _mock_factory):
+        from data import Data
+        import pandas as pd
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.resi = [1, 2]
+        d.out = pd.DataFrame({'x': range(5)})
+        cycle = d.resi_cycle()
+        self.assertEqual(len(cycle), 5)
+        self.assertEqual(cycle[:4], [1, 2, 1, 2])
+
+
+class TestStyleAndSave(unittest.TestCase):
+    "Table styling and CSV output"
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_style_creates_styler_and_csv(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.out = pd.DataFrame({'0': ['V2A'], 'test': [1.5]})
+            d.out_csv = os.path.join(tmpdir, 'out.csv')
+            d._style_and_save()
+            self.assertIsNotNone(d.out_table)
+            self.assertTrue(os.path.exists(d.out_csv))
+
+
+class TestParseSubEdgeCases(unittest.TestCase):
+    "parse_sub edge cases"
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_dms_duplicate_positions_allowed(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.seq = "MVEQYLL"
+        d.parse_sub("3 3")
+        # Duplicate positions produce duplicate mutation sets
+        self.assertEqual(len(d.resi), 2)
+        self.assertEqual(len(d.sub), 38)
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_mut_identical_sequences_no_muts(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.seq = "MV"
+        d.parse_sub("MV")  # identical — no differences
+        self.assertEqual(d.mode, 'MUT')
+        self.assertEqual(len(d.sub), 0)
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_single_residue_sequence_tms(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.model = FakeModel()
+        d.seq = "A"
+        d.parse_sub("scan all")
+        self.assertEqual(d.mode, 'DMS')
+        self.assertEqual(len(d.resi), 1)
+        self.assertEqual(len(d.sub), 19)
+
+
+class TestAppReturns(unittest.TestCase):
+    "app.py callback return values per mode"
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_image_property_returns_path_for_tms(self, _mock_factory):
+        from data import Data
+        d = object.__new__(Data)
+        d.out_table = None
+        d.out_img_path = 'out.png'
+        self.assertIsInstance(d.image, str)
+        self.assertEqual(d.image, 'out.png')
+
+    @patch('data.ModelFactory', return_value=FakeModel())
+    def test_image_property_returns_styler_for_dms(self, _mock_factory):
+        from data import Data
+        import pandas as pd
+        d = object.__new__(Data)
+        d.model_name = 'test'
+        d.out_csv = '/tmp/x.csv'
+        d.out = pd.DataFrame({'0': ['V2A'], 'test': [1.5]})
+        d._style_and_save()
+        # After styling, image returns Styler, not path
+        self.assertNotIsInstance(d.image, str)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/test_model.py

@@ -0,0 +1,109 @@
+import json
+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 TestModelFactoryRegistry(unittest.TestCase):
+    "ModelFactory register() and models()"
+
+    def setUp(self):
+        # Isolate _models for each test
+        from model import ModelFactory
+        self._orig = dict(ModelFactory._models)
+
+    def tearDown(self):
+        from model import ModelFactory
+        ModelFactory._models.clear()
+        ModelFactory._models.update(self._orig)
+
+    @patch('model.HfApi')
+    def test_register_adds_new_model(self, _mock_api):
+        from model import ModelFactory
+        class StubModel:
+            pass
+        ModelFactory.register("stub/model", StubModel)
+        self.assertIn("stub/model", ModelFactory.models())
+        self.assertEqual(ModelFactory._models["stub/model"], StubModel)
+
+    @patch('model.HfApi')
+    def test_models_returns_list_of_keys(self, _mock_api):
+        from model import ModelFactory
+        keys = ModelFactory.models()
+        self.assertIsInstance(keys, list)
+        self.assertTrue(len(keys) > 0)
+
+    @patch('model.HfApi')
+    def test_register_overwrites_existing(self, _mock_api):
+        from model import ModelFactory
+        existing_key = ModelFactory.models()[0]
+        class Replacement:
+            pass
+        ModelFactory.register(existing_key, Replacement)
+        self.assertIs(ModelFactory._models[existing_key], Replacement)
+
+
+class TestFetchCache(unittest.TestCase):
+    "_fetch_model_ids cache read/write/TTL"
+
+    @patch('model.time')
+    @patch('builtins.open', new_callable=unittest.mock.mock_open, read_data=json.dumps({"ts": 100, "ids": ["a/b"]}))
+    def test_cache_hit_within_ttl(self, mock_file, mock_time):
+        mock_time.return_value = 200  # 100s ago, within 3600 TTL
+        with patch('os.path.exists', return_value=True):
+            from model import _fetch_model_ids
+            result = _fetch_model_ids()
+            self.assertEqual(result, ["a/b"])
+
+    @patch('model.time')
+    @patch('builtins.open', new_callable=unittest.mock.mock_open, read_data=json.dumps({"ts": 100, "ids": ["a/b"]}))
+    def test_cache_miss_expired(self, mock_file, mock_time):
+        mock_time.return_value = 5000  # well past 3600 TTL
+        fallback_ids = ["fallback/model"]
+        mock_api = MagicMock()
+        mock_api.list_models.side_effect = [
+            [MagicMock(id="fallback/model")],
+            []
+        ]
+        with patch('os.path.exists', return_value=True), \
+             patch('model.HfApi', return_value=mock_api), \
+             tempfile.TemporaryDirectory() as tmpdir:
+            import model
+            orig_cache = model._MODEL_CACHE
+            model._MODEL_CACHE = os.path.join(tmpdir, "cache.json")
+            try:
+                result = model._fetch_model_ids()
+                self.assertIn("fallback/model", result)
+            finally:
+                model._MODEL_CACHE = orig_cache
+
+
+class TestScoringIndexExtraction(unittest.TestCase):
+    "Verify digit extraction from mutation strings (used in vectorized scoring)"
+
+    def _extract_idx(self, m):
+        """Replicate the logic from ESMModel.run_model"""
+        return int(''.join(c for c in m if c.isdigit())) - 1
+
+    def test_simple_two_digit(self):
+        self.assertEqual(self._extract_idx("V2A"), 1)
+
+    def test_three_digit_position(self):
+        self.assertEqual(self._extract_idx("R218K"), 217)
+
+    def test_single_digit_position(self):
+        self.assertEqual(self._extract_idx("M1D"), 0)
+
+    def test_various_positions(self):
+        cases = [("A10F", 9), ("L100W", 99), ("G500S", 499)]
+        for mut, expected in cases:
+            with self.subTest(mut=mut):
+                self.assertEqual(self._extract_idx(mut), expected)
+
+
+if __name__ == '__main__':
+    unittest.main()