164 lines
8.3 KiB
Python
164 lines
8.3 KiB
Python
import hashlib
|
|
import importlib.util
|
|
import io
|
|
import json
|
|
from contextlib import redirect_stdout
|
|
from pathlib import Path
|
|
import py_compile
|
|
import subprocess
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent.parent
|
|
SPEC = importlib.util.spec_from_file_location('semantic_retrieval', ROOT / 'scripts/lib/semantic-retrieval.py')
|
|
MODULE = importlib.util.module_from_spec(SPEC)
|
|
SPEC.loader.exec_module(MODULE)
|
|
|
|
|
|
class SemanticRetrievalTests(unittest.TestCase):
|
|
def test_sync_does_not_register_mcp(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
home = Path(directory) / 'home'
|
|
output = io.StringIO()
|
|
with redirect_stdout(output):
|
|
MODULE.sync(ROOT, home, ['codex', 'claude'], True, True, False)
|
|
self.assertNotIn('mcp', output.getvalue().lower())
|
|
|
|
def test_runtime_prefers_cuda_packages_for_nvidia_gpu(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
target = Path(directory)
|
|
python = MODULE.python_path(target / '.venv')
|
|
python.parent.mkdir(parents=True)
|
|
python.touch()
|
|
with patch.object(MODULE.shutil, 'which', return_value='nvidia-smi'), patch.object(MODULE.subprocess, 'run') as run:
|
|
MODULE.ensure_runtime(target, False)
|
|
|
|
self.assertEqual(run.call_args.args[0][-2:], ['--extra-index-url', MODULE.CUDA_INDEX])
|
|
|
|
def test_runtime_falls_back_to_cpu_packages_when_cuda_install_fails(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
target = Path(directory)
|
|
python = MODULE.python_path(target / '.venv')
|
|
python.parent.mkdir(parents=True)
|
|
python.touch()
|
|
failure = MODULE.subprocess.CalledProcessError(1, 'pip')
|
|
with patch.object(MODULE.shutil, 'which', return_value='nvidia-smi'), patch.object(MODULE.subprocess, 'run', side_effect=[failure, None]) as run:
|
|
MODULE.ensure_runtime(target, False)
|
|
|
|
self.assertEqual(run.call_count, 2)
|
|
self.assertNotIn('--extra-index-url', run.call_args.args[0])
|
|
|
|
def test_sync_preloads_models_after_installing_runtime(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
home = Path(directory) / 'home'
|
|
target = home / '.my-ai-configuration/semantic-retrieval'
|
|
python = MODULE.python_path(target / '.venv')
|
|
with patch.object(MODULE, 'ensure_runtime', return_value=python), \
|
|
patch.object(MODULE.subprocess, 'run') as run:
|
|
MODULE.sync(ROOT, home, ['codex'], True, False, False)
|
|
|
|
self.assertEqual([[
|
|
str(python), str(target / 'server.py'), 'preload', '--model-cache', str(target / 'model-cache')
|
|
]], [call.args[0] for call in run.call_args_list])
|
|
|
|
def test_summary_dry_run_omits_details_without_writes(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
home = Path(directory) / 'home'
|
|
for enabled in (True, False):
|
|
output = io.StringIO()
|
|
with redirect_stdout(output):
|
|
MODULE.sync(ROOT, home, ['codex'], enabled, True, False, summary=True)
|
|
self.assertEqual(output.getvalue(), '')
|
|
self.assertFalse(home.exists())
|
|
|
|
def test_disable_removes_owned_runtime_without_cli_registration(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
home = Path(directory) / 'home'
|
|
target = home / '.my-ai-configuration/semantic-retrieval'
|
|
target.mkdir(parents=True)
|
|
for name in MODULE.FILES:
|
|
(target / name).write_text(name)
|
|
state_path = target.parent / 'semantic-retrieval.json'
|
|
state_path.write_text(json.dumps({'version': 1, 'clients': ['codex'], 'files': {
|
|
name: hashlib.sha256((target / name).read_bytes()).hexdigest() for name in MODULE.FILES}}))
|
|
MODULE.sync(ROOT, home, ['codex'], False, False, False)
|
|
self.assertFalse(target.exists())
|
|
self.assertFalse(state_path.exists())
|
|
|
|
def test_update_accepts_python_cache_but_rejects_foreign_cache_content(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
home = Path(directory) / 'home'
|
|
target = home / '.my-ai-configuration/semantic-retrieval'
|
|
target.mkdir(parents=True)
|
|
files = {}
|
|
for name in MODULE.FILES:
|
|
content = (ROOT / 'shared/retrieval' / name).read_bytes()
|
|
(target / name).write_bytes(content)
|
|
files[name] = hashlib.sha256(content).hexdigest()
|
|
MODULE.save_state(target.parent / 'semantic-retrieval.json', {'version': 1, 'clients': ['codex'], 'files': files}, False)
|
|
py_compile.compile(str(target / 'server.py'), doraise=True)
|
|
cache = target / '__pycache__'
|
|
with patch.object(MODULE, 'ensure_runtime', return_value=Path('python')), patch.object(MODULE, 'ensure_models'):
|
|
MODULE.sync(ROOT, home, ['codex'], True, False, True)
|
|
(cache / 'foreign.txt').write_text('keep')
|
|
with self.assertRaisesRegex(ValueError, '__pycache__'):
|
|
MODULE.sync(ROOT, home, ['codex'], True, False, True)
|
|
|
|
def test_install_search_update_and_remove_owned_data(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
home = Path(directory) / 'home'
|
|
repository = Path(directory) / 'repository'
|
|
repository.mkdir()
|
|
subprocess.run(['git', '-C', str(repository), 'init', '--quiet'], check=True)
|
|
target = home / '.my-ai-configuration/semantic-retrieval'
|
|
state_path = target.parent / 'semantic-retrieval.json'
|
|
with patch.object(MODULE, 'ensure_runtime', return_value=Path('python')), \
|
|
patch.object(MODULE, 'ensure_models'):
|
|
MODULE.sync(ROOT, home, ['codex'], True, False, False)
|
|
installed = importlib.util.spec_from_file_location('installed_retrieval', target / 'server.py')
|
|
server = importlib.util.module_from_spec(installed)
|
|
installed.loader.exec_module(server)
|
|
index = server.Index(repository, target / 'data')
|
|
try:
|
|
self.assertEqual(index.search('query'), [])
|
|
finally:
|
|
index.db.close()
|
|
self.assertEqual(len(list((target / 'data').glob('*.sqlite3'))), 1)
|
|
server.daemon_authkey(target / 'data')
|
|
MODULE.sync(ROOT, home, ['codex'], True, False, True)
|
|
self.assertTrue((target / 'data').is_dir())
|
|
foreign = target / 'data' / 'foreign.txt'
|
|
foreign.write_text('keep')
|
|
with self.assertRaisesRegex(ValueError, 'data'):
|
|
MODULE.sync(ROOT, home, ['codex'], False, False, False)
|
|
self.assertTrue(foreign.exists())
|
|
foreign.unlink()
|
|
MODULE.sync(ROOT, home, ['codex'], False, False, False)
|
|
self.assertFalse(target.exists())
|
|
self.assertFalse(state_path.exists())
|
|
|
|
def test_negative_benchmark_keeps_active_runtime(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
home = Path(directory) / 'home'
|
|
target = home / '.my-ai-configuration/semantic-retrieval'
|
|
target.mkdir(parents=True)
|
|
state_path = target.parent / 'semantic-retrieval.json'
|
|
files = {}
|
|
for name in MODULE.FILES:
|
|
content = (ROOT / 'shared/retrieval' / name).read_bytes()
|
|
(target / name).write_bytes(content)
|
|
files[name] = hashlib.sha256(content).hexdigest()
|
|
MODULE.save_state(state_path, {'version': 1, 'clients': ['codex'], 'files': files}, False)
|
|
result = type('Result', (), {'stdout': '{"recommended": false, "embeddingSeconds": 1, "rerankingSeconds": 1}'})()
|
|
with patch.object(MODULE, 'ensure_runtime', return_value=Path('python')), \
|
|
patch.object(MODULE.subprocess, 'run', return_value=result):
|
|
self.assertFalse(MODULE.benchmark(ROOT, home, False, True))
|
|
self.assertTrue(target.is_dir())
|
|
self.assertEqual(MODULE.load_state(state_path)['clients'], ['codex'])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|