Files
nest/tests/test_tool_jobs.py
T
kstyagi 3e17e3e72f feat(tools): Enhance guest tools installation and management
- Added new API endpoints for tool catalog and installation management.
- Updated usage documentation to include new tool installation features.
- Implemented asynchronous tool installation with job tracking.
- Introduced integration tests for tool installation in disposable VMs.
- Enhanced dashboard UI to reflect tool installation status and usage instructions.
- Improved error handling and logging for tool installation processes.
- Added support for multiple Python versions and Conda environments.
- Refactored tool management logic to ensure better state handling and job persistence.
2026-09-25 23:52:21 +05:30

150 lines
7.2 KiB
Python

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"])