kg-agent/scripts/test_bootstrap.py

"""Bootstrap sequencing with a fake uv; no installs, downloads or services."""
import json
import os
from pathlib import Path
import shutil
import subprocess
import tempfile
import unittest
from unittest.mock import patch
from scripts import download_model

ROOT = Path(__file__).resolve().parents[1]
BASH = shutil.which('bash') or ('C:/Program Files/Git/bin/bash.exe' if Path('C:/Program Files/Git/bin/bash.exe').exists() else None)


@unittest.skipUnless(BASH, 'bash is unavailable')
class BootstrapTests(unittest.TestCase):
    def run_script(self, action, fail='', env_exists=True, args=()):
        with tempfile.TemporaryDirectory() as tmp:
            root=Path(tmp)
            shutil.copy2(ROOT/'deploy.sh',root/'deploy.sh')
            (root/'.env.example').write_text('EXAMPLE=1\n')
            if env_exists: (root/'.env').write_text('KEEP=unchanged\n')
            bin_dir=root/'bin'; bin_dir.mkdir()
            fake=bin_dir/'uv'
            fake.write_text('#!/bin/sh\nprintf "%s\\n" "$*" >> "$CALL_LOG"\ncase "$*" in *"$FAIL_MATCH"*) if [ -n "$FAIL_MATCH" ]; then exit 9; fi ;; esac\n')
            fake.chmod(0o755)
            env={**os.environ,'PATH':str(bin_dir)+os.pathsep+os.environ['PATH'],'CALL_LOG':str(root/'calls.log'),'FAIL_MATCH':fail}
            result=subprocess.run([BASH,str(root/'deploy.sh'),action,*args],env=env,capture_output=True,text=True)
            calls=(root/'calls.log').read_text() if (root/'calls.log').exists() else ''
            return result.returncode,calls,(root/'.env').read_text()

    def test_templates_and_dms_bootstrap_then_deploy(self):
        for action in ('templates','dms','start'):
            code,calls,env=self.run_script(action)
            self.assertEqual(code,0,calls)
            self.assertLess(calls.index('python install'),calls.index('sync --frozen'))
            self.assertLess(calls.index('sync --frozen'),calls.index('download_model.py'))
            lines=calls.splitlines()
            self.assertTrue(lines[-1].endswith('scripts/start_service.py'))
            if action == 'start':
                self.assertNotIn('update_data.py',calls)
            else:
                self.assertIn('update_data.py --mode '+action,lines[-2])
                self.assertLess(calls.index('download_model.py'),calls.index('update_data.py'))
            self.assertEqual(env,'KEEP=unchanged\n')

    def test_restart_starts_with_restart_flag_without_data_update(self):
        code,calls,_=self.run_script('restart')
        self.assertEqual(code,0)
        self.assertNotIn('update_data.py',calls)
        self.assertTrue(calls.splitlines()[-1].endswith('start_service.py --restart'))

    def test_setup_never_runs_business_deployment(self):
        code,calls,_=self.run_script('setup')
        self.assertEqual(code,0)
        self.assertIn('download_model.py',calls)
        self.assertNotIn('update_data.py',calls)
        self.assertNotIn('start_service.py',calls)

    def test_environment_and_download_failures_stop_pipeline(self):
        for failure in ('sync --frozen','download_model.py'):
            code,calls,_=self.run_script('templates',failure)
            self.assertEqual(code,9)
            self.assertNotIn('update_data.py',calls)
            self.assertNotIn('scripts/start_service.py\n',calls)

    def test_update_failure_prevents_start(self):
        for action in ('templates','dms'):
            code,calls,_=self.run_script(action,'update_data.py')
            self.assertEqual(code,9)
            self.assertIn('update_data.py --mode '+action,calls)
            self.assertNotIn('scripts/start_service.py\n',calls)

    def test_start_failure_and_arguments(self):
        code,calls,_=self.run_script('start',args=('--host','127.0.0.1','--port','8123'))
        self.assertEqual(code,0)
        self.assertTrue(calls.splitlines()[-1].endswith('start_service.py --host 127.0.0.1 --port 8123'))
        code,_,_=self.run_script('start','start_service.py --host',args=('--host','127.0.0.1'))
        self.assertEqual(code,9)

    def test_invalid_config_prevents_model_and_update(self):
        code,calls,_=self.run_script('templates','--check-config')
        self.assertEqual(code,9)
        self.assertNotIn('download_model.py',calls)
        self.assertNotIn('update_data.py',calls)

    def test_missing_env_is_created_without_business_update(self):
        code,calls,env=self.run_script('templates',env_exists=False)
        self.assertEqual(code,2)
        self.assertEqual(calls,'')
        self.assertEqual(env,'EXAMPLE=1\n')


class ModelTests(unittest.TestCase):
    def test_incomplete_shards_do_not_count_as_ready(self):
        with tempfile.TemporaryDirectory() as tmp:
            root=Path(tmp)
            for name in ('config.json','tokenizer.json','tokenizer_config.json'):
                (root/name).write_text('{}')
            (root/'modules.json').write_text('[{"type":"sentence_transformers.models.Normalize","path":"2_Normalize"}]')
            (root/'model.safetensors.index.json').write_text(json.dumps({'weight_map':{'a':'part1','b':'part2'}}))
            (root/'part1').write_text('weight')
            self.assertFalse(download_model.model_files_ready(root))
            (root/'part2').write_text('weight')
            self.assertTrue(download_model.model_files_ready(root))

    def test_ready_model_skips_network_but_verifies_loading(self):
        with patch.object(download_model,'get_embedding_model_dir',return_value=Path('model')), patch.object(download_model,'model_files_ready',return_value=True), patch.object(download_model,'verify_model') as verify, patch.object(download_model.subprocess,'run') as run:
            download_model.ensure_model()
            run.assert_not_called()
            verify.assert_called_once()

    def test_check_only_missing_model_never_downloads(self):
        with patch.object(download_model,'get_embedding_model_dir',return_value=Path('model')), patch.object(download_model,'model_files_ready',return_value=False), patch.object(download_model.subprocess,'run') as run:
            with self.assertRaises(RuntimeError): download_model.ensure_model(check_only=True)
            run.assert_not_called()


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