kstyagi
e1b934e35e
- Added `guest_tools.py` for standalone tool installer functionality. - Introduced `ToolJobStore` in `tool_store.py` to manage per-VM tool job records. - Created `ToolService` in `tools.py` for background installation workers using the guest execution channel. - Enhanced `LifecycleService` in `lifecycle.py` to support tool installation queuing and preparation. - Updated `guest_exec.py` to handle connection shutdown gracefully. - Added validation and normalization functions for tool specifications. - Implemented logging and error handling for tool installation processes. - Ensured compatibility with Ubuntu 22.04 and ARM64/x86_64 architectures.
152 lines
7.2 KiB
Python
152 lines
7.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import threading
|
|
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"])
|