moebiusT7/book-ocr-studio

1import importlib.util2import json3import os4import subprocess5import sys6import tempfile7import unittest8from pathlib import Path9from unittest.mock import MagicMock, patch10 11ROOT = Path(__file__).resolve().parents[1]12sys.path.insert(0, str(ROOT))13from model_store import DEFAULT_MODEL, OPTIONAL_MODEL, model_directory14spec = importlib.util.spec_from_file_location('setup_models', ROOT/'scripts/setup_models.py')15setup = importlib.util.module_from_spec(spec)16spec.loader.exec_module(setup)17 18class SetupTests(unittest.TestCase):19    def exercise(self, initial, final, *, failure=False):20        with tempfile.TemporaryDirectory() as tmp, patch.dict(os.environ, {}, clear=False):21            os.environ.pop('BOOK_OCR_MODELS', None)22            root = Path(tmp)23            client = MagicMock()24            client.get.return_value.json.side_effect = [{'models': initial}, {'models': final}]25            with patch.object(setup.requests, 'Session') as session, patch.object(setup.subprocess, 'Popen') as popen, patch.object(setup.subprocess, 'run') as pull:26                session.return_value.__enter__.return_value = client27                popen.return_value.poll.return_value = None28                if failure:29                    pull.side_effect = subprocess.CalledProcessError(1, 'synthetic-pull')30                    with self.assertRaises(subprocess.CalledProcessError):31                        setup.prepare(root, DEFAULT_MODEL, executable='synthetic-ollama')32                    self.assertFalse((root/'model-settings.json').exists())33                elif not final:34                    with self.assertRaises(RuntimeError):35                        setup.prepare(root, DEFAULT_MODEL, executable='synthetic-ollama')36                    self.assertFalse((root/'model-settings.json').exists())37                else:38                    setup.prepare(root, DEFAULT_MODEL, executable='synthetic-ollama')39                    data = json.loads((root/'model-settings.json').read_text())40                    self.assertEqual(data['models'][DEFAULT_MODEL], 'synthetic-digest')41                    self.assertEqual(model_directory(root), root/'models')42                    self.assertEqual((root/'model-settings.json').stat().st_mode & 0o777, 0o600)43                if initial:44                    pull.assert_not_called()45                else:46                    self.assertEqual(pull.call_args.args[0], ['synthetic-ollama', 'pull', DEFAULT_MODEL])47                self.assertEqual(popen.call_args.kwargs['env']['OLLAMA_MODELS'], str(root/'models'))48                self.assertTrue(popen.call_args.kwargs['env']['OLLAMA_HOST'].startswith('127.0.0.1:'))49                popen.return_value.terminate.assert_called_once()50 51    def test_default_download_and_record(self):52        self.exercise([], [{'name': DEFAULT_MODEL, 'digest': 'synthetic-digest'}])53 54    def test_existing_model_reused(self):55        model = {'name': DEFAULT_MODEL, 'digest': 'synthetic-digest'}56        self.exercise([model], [model])57 58    def test_failed_download_has_no_success_record(self):59        self.exercise([], [], failure=True)60 61    def test_missing_model_after_pull_has_no_success_record(self):62        self.exercise([], [])63 64    def test_corrupt_configuration_does_not_silently_fallback(self):65        with tempfile.TemporaryDirectory() as tmp, patch.dict(os.environ, {}, clear=False):66            os.environ.pop('BOOK_OCR_MODELS', None)67            root = Path(tmp)68            (root/'model-settings.json').write_text('{}')69            with self.assertRaises(ValueError):70                model_directory(root)71            os.environ['BOOK_OCR_MODELS'] = str(root/'override')72            self.assertEqual(model_directory(root), root/'override')73 74    def test_optional_model_does_not_change_default(self):75        self.assertEqual(DEFAULT_MODEL, 'gemma4:12b-it-qat')76        self.assertEqual(OPTIONAL_MODEL, 'gemma4:26b-a4b-it-qat')77 78if __name__ == '__main__':79    unittest.main()80