import hashlib import importlib.util import io import json from contextlib import redirect_stdout from pathlib import Path import tempfile import unittest from unittest.mock import patch ROOT = Path(__file__).resolve().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_codex_registration_sets_tool_timeout(self): with tempfile.TemporaryDirectory() as directory: config = Path(directory) / 'config.toml' config.write_text('[mcp_servers.my-ai-qwen3-retrieval]\ncommand = "python"\n\n[mcp_servers.my-ai-qwen3-retrieval.env]\nHF_HOME = "cache"\n', encoding='utf-8') MODULE.set_codex_tool_timeout(config, 1800) self.assertEqual(config.read_text(encoding='utf-8'), '[mcp_servers.my-ai-qwen3-retrieval]\ncommand = "python"\ntool_timeout_sec = 1800\n\n[mcp_servers.my-ai-qwen3-retrieval.env]\nHF_HOME = "cache"\n') 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_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_claude_registration_places_name_before_environment(self): with tempfile.TemporaryDirectory() as directory: home = Path(directory) / 'home' calls = [] original = MODULE.command original_runtime = MODULE.ensure_runtime MODULE.command = lambda client, arguments, home, dry_run: calls.append((client, arguments)) MODULE.ensure_runtime = lambda target, dry_run: target / '.venv' / 'bin' / 'python' try: MODULE.sync(ROOT, home, ['claude'], True, False, False) finally: MODULE.command = original MODULE.ensure_runtime = original_runtime add = calls[0][1] self.assertEqual(add[:5], ['claude', 'mcp', 'add', MODULE.NAME, '--scope']) self.assertIn('-e', add) def test_disable_removes_owned_runtime_and_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}})) calls = [] original = MODULE.command MODULE.command = lambda client, arguments, home, dry_run: calls.append((client, arguments)) try: MODULE.sync(ROOT, home, ['codex'], False, False, False) finally: MODULE.command = original self.assertEqual(calls[0][1], ['codex', 'mcp', 'remove', MODULE.NAME]) self.assertFalse(target.exists()) self.assertFalse(state_path.exists()) if __name__ == '__main__': unittest.main()