from __future__ import annotations import json import tempfile import unittest from concurrent.futures import ThreadPoolExecutor from dataclasses import replace from pathlib import Path from unittest.mock import Mock, patch from tests.test_server import FakeLifecycle from uvm.config import Settings from uvm.errors import ConflictError, GuestExecutionError, NotFoundError, StateError from uvm.state import StateStore from uvm.tool_catalog import ToolSelection from uvm.tool_store import ToolJobStore from uvm.tools import ToolService class ToolJobTests(unittest.TestCase): def setUp(self): temporary = self.enterContext(tempfile.TemporaryDirectory()) self.settings = replace(Settings(), base=Path(temporary), tool_workers=1) self.state = StateStore(self.settings) self.vm = replace(FakeLifecycle().vm, runtime_owned=True, pid=123, process_start_time="1", boot_id="host-boot") self.settings.vm_dir(self.vm.id).mkdir(parents=True, mode=0o700) with self.state.transaction() as state: state.vms[self.vm.id] = self.vm self.store = ToolJobStore(self.settings, self.state) def enqueue(self, selections=None, **options): with self.state.operation_lock(): return self.store.enqueue_locked(self.vm.id, selections or [ToolSelection("git")], **options) def test_duplicate_concurrent_submissions_have_one_durable_job(self): with ThreadPoolExecutor(max_workers=4) as executor: jobs = list(executor.map(lambda _: self.enqueue(request_id="request-1"), range(4))) self.assertEqual(len({job["id"] for job in jobs}), 1) self.assertEqual(len(self.store.read(self.vm.id)["jobs"]), 1) path = self.settings.vm_dir(self.vm.id) / "tools.json" self.assertEqual(path.stat().st_mode & 0o777, 0o600) self.assertNotIn("guest-secret", path.read_text()) with self.assertRaises(ConflictError): self.enqueue([ToolSelection("htop")]) def test_history_is_vm_scoped_and_deletion_is_never_undone(self): job = self.enqueue() with self.assertRaises(NotFoundError): self.store.get_job(self.vm.name, job["id"]) with self.state.transaction() as state: del state.vms[self.vm.id] with self.assertRaises(NotFoundError): self.store.update_job(self.vm.id, job["id"], lambda current, data: current.update(status="succeeded")) self.assertNotIn(self.vm.id, self.state.load().vms) def test_recovery_preserves_queued_work_and_marks_inflight_outcomes_unknown(self): job = self.enqueue() self.store.recover() self.assertEqual(self.store.get_job(self.vm.id, job["id"])["status"], "queued") self.store.claim() self.store.update_job(self.vm.id, job["id"], lambda current, data: current["items"][0].update(status="installing")) self.store.recover() recovered = self.store.get_job(self.vm.id, job["id"]) self.assertEqual(recovered["status"], "interrupted") self.assertTrue(recovered["outcome_unknown"]) self.assertEqual(recovered["items"][0]["status"], "interrupted") self.assertIsNone(self.store.claim()) def test_symlinks_and_foreign_history_are_rejected(self): outside = self.settings.base / "outside" outside.write_text("private data") path = self.settings.vm_dir(self.vm.id) / "tools.json" path.symlink_to(outside) with self.assertRaises((StateError, OSError)): self.store.read(self.vm.id) self.assertEqual(outside.read_text(), "private data") path.unlink() self.enqueue() data = json.loads(path.read_text()) data["jobs"][0]["vm_id"] = "vm-other" path.write_text(json.dumps(data)) with self.assertRaises(StateError): self.store.read(self.vm.id) def test_only_one_worker_can_own_this_registry_and_shutdown_releases_it(self): first = ToolService(self.settings, self.state, Mock()) second = ToolService(self.settings, self.state, Mock()) self.addCleanup(first.shutdown) self.addCleanup(second.shutdown) first.start() with self.assertRaises(ConflictError): second.start() first.shutdown() second.start() def make_service(self): lifecycle = Mock() lifecycle.prepare_tool_install.return_value = self.vm service = ToolService(self.settings, self.state, lifecycle) return service @staticmethod def observation(selection, installed=False): return {**selection.to_dict(), "key": selection.key, "installed": installed, "detected_version": "2.0.0" if installed else None, "usage": "tool --version", "error": None} def test_worker_reports_partial_failure_and_retains_successful_tools(self): selections = [ToolSelection("git"), ToolSelection("htop")] job = self.enqueue(selections) claimed = self.store.claim() service = self.make_service() results = [ {"ok": True, "busy": False, "boot_id": "guest-boot", "tools": [self.observation(item) for item in selections]}, {"ok": True, "status": "installed", "tool": self.observation(selections[0], True), "output": "git installed"}, {"ok": False, "error": "mirror unavailable", "output": "apt download failed"}, ] with patch.object(service, "_execute", side_effect=results) as execute: service._run_job(claimed) result = self.store.get_job(self.vm.id, job["id"]) self.assertEqual(result["status"], "failed") self.assertEqual([item["status"] for item in result["items"]], ["installed", "failed"]) self.assertEqual(result["items"][1]["output"], "apt download failed") self.assertEqual(self.state.get(self.vm.id).status, "running") self.assertEqual(execute.call_args.args[1]["boot_id"], "guest-boot") def test_lost_result_is_never_replayed_and_remaining_tools_are_skipped(self): selections = [ToolSelection("git"), ToolSelection("htop")] job = self.enqueue(selections) service = self.make_service() with patch.object(service, "_execute", side_effect=[ {"ok": True, "busy": False, "boot_id": "boot", "tools": [self.observation(item) for item in selections]}, GuestExecutionError("connection lost", outcome_unknown=True), ]) as execute: service._run_job(self.store.claim()) result = self.store.get_job(self.vm.id, job["id"]) self.assertEqual(execute.call_count, 2) self.assertEqual(result["status"], "interrupted") self.assertTrue(result["outcome_unknown"]) self.assertEqual([item["status"] for item in result["items"]], ["interrupted", "skipped"]) def test_prior_guest_installation_blocks_a_retry_before_submission(self): job = self.enqueue() service = self.make_service() with patch.object(service, "_execute", return_value={"ok": True, "busy": True}) as execute: service._run_job(self.store.claim()) self.assertEqual(execute.call_count, 1) result = self.store.get_job(self.vm.id, job["id"]) self.assertEqual(result["status"], "interrupted") self.assertFalse(result["outcome_unknown"])