diff --git a/python/packages/jumpstarter-driver-cuttlefish/README.md b/python/packages/jumpstarter-driver-cuttlefish/README.md index 5f4050108..94ad7e8ec 100644 --- a/python/packages/jumpstarter-driver-cuttlefish/README.md +++ b/python/packages/jumpstarter-driver-cuttlefish/README.md @@ -173,7 +173,7 @@ export: This is a **composite driver** with three children: - **power** — `VirtualPowerInterface`: `j power on`, `j power off [--destroy]`, `j power cycle` -- **storage** — `FlasherInterface`: not yet implemented (planned: HO artifact upload API) +- **storage** — `FlasherInterface`: flashes Cuttlefish image archives through the HO user-artifact API - **adb** — ADB server for device communication The exporter config also typically includes sibling drivers: @@ -182,6 +182,53 @@ The exporter config also typically includes sibling drivers: Use `ref:` entries in the exporter config to expose children at the top level. +### Storage flash + +Storage flash uploads archives to Host Orchestrator, extracts them, and merges +their contents into one image directory. The documented target names are: + +- `images` — the device image ZIP (`*-img-*.zip`) +- `host_package` — the Cuttlefish host package (`cvd-host_package.tar.gz`) + +For a split AAOS build, flash both artifacts in one call: + +```console +j storage flash \ + --target images:/path/to/aosp_cf_x86_64_auto-img-.zip \ + --target host_package:/path/to/cvd-host_package.tar.gz +``` + +The same targets accept `http://` or `https://` URLs. The client passes those +URLs as presigned GET resources; the exporter downloads and streams them into +the same hashing, deduplication, extraction, and staging path, so the artifact +does not pass through the client. + +An untargeted single archive is treated as a complete bundle and is injected +into both `common.host_package` and every +`instances[].disk.default_build`. A targeted call may provide only one role; +when an active generation exists, the missing role is carried forward by +checksum into a new immutable image directory. This means +`flash(host_package)` after a complete flash keeps the existing device images. +For a fresh CVD, provide both targeted artifacts (or a complete untargeted +bundle). Targeted roles must match the archive format; an unknown target is +rejected. + +Flashing only stages the files. Recreate the CVD to select them: + +```console +j power off --destroy +j power on +``` + +The staged and active image-directory generations are exporter in-memory state. +An exporter restart loses that state; the user artifacts remain in Host +Orchestrator, but the configured `env_config` is used until the artifacts are +flashed again. `power off --destroy` releases the active image directory after +the CVD is deleted. At exporter teardown, pending staged directories are +deleted best-effort; an active directory is retained because its CVD may still +be running. Cleanup failures are logged at WARNING and may require manual +cleanup on the Host Orchestrator. + ## Usage ### CLI diff --git a/python/packages/jumpstarter-driver-cuttlefish/jumpstarter_driver_cuttlefish/driver.py b/python/packages/jumpstarter-driver-cuttlefish/jumpstarter_driver_cuttlefish/driver.py index 42ec00deb..833da3fa7 100644 --- a/python/packages/jumpstarter-driver-cuttlefish/jumpstarter_driver_cuttlefish/driver.py +++ b/python/packages/jumpstarter-driver-cuttlefish/jumpstarter_driver_cuttlefish/driver.py @@ -1,10 +1,17 @@ +import copy +import hashlib import json +import os import subprocess +import tempfile import time from collections.abc import Generator from dataclasses import dataclass, field +import httpx import requests +from anyio import EndOfStream, sleep, to_thread +from anyio.streams.file import FileReadStream, FileWriteStream from jumpstarter_driver_adb.driver import AdbServer from jumpstarter_driver_power.driver import PowerReading, VirtualPowerInterface @@ -20,6 +27,14 @@ class CuttlefishTimeout(CuttlefishError): """Raised when an operation doesn't complete in time.""" +@dataclass +class ImageDirGeneration: + """A Cuttlefish image directory and the artifacts populated into it.""" + + dir_id: str + artifacts: dict[str, str] = field(default_factory=dict) + + @dataclass(kw_only=True) class Cuttlefish(Driver): """Cuttlefish Host Orchestrator driver for managing Android virtual devices. @@ -39,8 +54,15 @@ class Cuttlefish(Driver): boot_timeout: int = 300 env_config: dict = field(default_factory=dict) webrtc_url: str = "" + upload_chunk_size: int = 16 * 1024 * 1024 + upload_max_retries: int = 3 _cvd_group: str | None = field(default=None, init=False, repr=False) _cvd_name: str | None = field(default=None, init=False, repr=False) + # Both generations are in-memory session state. The active directory remains + # immutable while a CVD references it; a partial flash creates a new + # generation and repopulates it from the active checksums. + _staged: ImageDirGeneration | None = field(default=None, init=False, repr=False) + _active: ImageDirGeneration | None = field(default=None, init=False, repr=False) def __post_init__(self): if hasattr(super(), "__post_init__"): @@ -68,22 +90,43 @@ def _cvd_path(self) -> str: def _fmt(self, result) -> str: return json.dumps(result, indent=2) if isinstance(result, (dict, list)) else str(result) - def _request(self, method: str, path: str, data: dict | None = None, timeout: float = 10) -> dict | list | str: + def _request_response( + self, method: str, path: str, data: dict | None = None, timeout: float = 10 + ) -> requests.Response: try: - r = requests.request(method, f"{self._base_url}{path}", json=data, timeout=timeout) - r.raise_for_status() - try: - return r.json() - except requests.JSONDecodeError: - return r.text + return requests.request(method, f"{self._base_url}{path}", json=data, timeout=timeout) except requests.ConnectionError as e: raise CuttlefishError(f"not connected to Host Orchestrator at {self.host}:{self.port}") from e except requests.Timeout as e: raise CuttlefishError(f"{method} {path} timed out after {timeout}s") from e + + def _raise_for_status(self, response: requests.Response, context: str) -> None: + try: + response.raise_for_status() except requests.HTTPError as e: - raise CuttlefishError(f"{method} {path} failed: {e}") from e + raise CuttlefishError(f"{context} failed: {e}") from e - def _wait_for_operation(self, op_name: str, timeout: float = 300) -> dict: + def _response_body(self, response: requests.Response) -> dict | list | str | None: + if not response.content: + return None + try: + return response.json() + except requests.JSONDecodeError: + return response.text + + def _request( + self, method: str, path: str, data: dict | None = None, timeout: float = 10 + ) -> dict | list | str | None: + response = self._request_response(method, path, data, timeout) + self._raise_for_status(response, f"{method} {path}") + return self._response_body(response) + + def _wait_for_operation( + self, + op_name: str, + timeout: float = 300, + accepted_statuses: tuple[int, ...] = (), + ) -> dict | list | str | None: deadline = time.monotonic() + timeout start = time.monotonic() while True: @@ -107,31 +150,35 @@ def _wait_for_operation(self, op_name: str, timeout: float = 300) -> dict: self.logger.info("operation %s: server busy (%d), retrying in 2s", op_name, r.status_code) time.sleep(2) continue - if r.status_code == 500: - body = None - try: - body = r.json() - except (ValueError, requests.JSONDecodeError): - pass - if body and isinstance(body, dict): - msg = body.get("error", "unknown error") - details = body.get("details", "") - raise CuttlefishError(f"operation failed: {msg}\n{details}") - raise CuttlefishError(f"operation failed with status 500: {r.text}") - try: - r.raise_for_status() - except requests.HTTPError as e: - raise CuttlefishError(f"operation {op_name} failed: {e}") from e - return r.json() + return self._operation_result(op_name, r, accepted_statuses) raise CuttlefishTimeout(f"operation {op_name} timed out after {timeout}s") + def _operation_result( + self, + op_name: str, + response: requests.Response, + accepted_statuses: tuple[int, ...], + ) -> dict | list | str | None: + if response.status_code == 500: + body = self._response_body(response) + if isinstance(body, dict): + msg = body.get("error", "unknown error") + details = body.get("details", "") + raise CuttlefishError(f"operation failed: {msg}\n{details}") + raise CuttlefishError(f"operation failed with status 500: {response.text}") + if response.status_code in accepted_statuses: + self.logger.info("operation %s completed with accepted status %d", op_name, response.status_code) + return None + self._raise_for_status(response, f"operation {op_name}") + return self._response_body(response) + def _do_operation( self, method: str, path: str, data: dict | None = None, timeout: float = 300, - ) -> dict | list | str: + ) -> dict | list | str | None: result = self._request(method, path, data) if isinstance(result, dict) and "done" in result: op_name = result.get("name") @@ -141,6 +188,145 @@ def _do_operation( return self._wait_for_operation(str(op_name), timeout) return result + def _artifact_exists(self, checksum: str) -> bool: + """Return True if a user artifact with this SHA-256 is already uploaded. + + HO addresses artifacts by checksum, so an existing upload can be reused + across sessions and hosts — a re-flash of the same image skips the + (potentially multi-GB) transfer. + """ + r = self._request_response("GET", f"/v1/userartifacts/{checksum}") + if r.status_code == 404: + return False + self._raise_for_status(r, "stat user artifact") + return True + + def _extract_artifact(self, checksum: str) -> None: + """Extract an uploaded artifact server-side (async operation). + + HO dispatches on the *stored file's suffix* (``.zip`` → unzip, + ``.tar.gz`` → untar), so the upload must have preserved the real + extension. Extracting an already-extracted checksum returns 409 + Conflict, which we treat as success — this is the idempotent re-flash + path (upload skipped because the artifact was already present). + """ + r = self._request_response("POST", f"/v1/userartifacts/{checksum}/:extract", timeout=130) + if r.status_code == 409: + self.logger.info("artifact %s already extracted", checksum[:12]) + return + self._raise_for_status(r, "extract artifact") + result = self._response_body(r) + if result is None: + return + if isinstance(result, dict) and "done" in result: + op_name = result.get("name") + if not op_name: + raise CuttlefishError(f"extract operation missing 'name': {result}") + self._wait_for_operation(str(op_name), accepted_statuses=(409,)) + + def _delete_image_dir_quietly(self, dir_id: str) -> bool: + """Best-effort DELETE of an image dir; log but don't raise on failure. + + HO refuses to delete a dir a running CVD still references, so callers + must be able to tolerate failure (e.g. teardown while a CVD is up). + """ + try: + self._do_operation("DELETE", f"/cvd_imgs_dirs/{dir_id}") + except CuttlefishError as e: + self.logger.warning("failed to delete image dir %s; it may need manual cleanup: %s", dir_id, e) + return False + return True + + def _create_image_dir(self) -> str: + """Create an empty image directory, returning its id (async operation).""" + result = self._do_operation("POST", "/cvd_imgs_dirs") + dir_id = None + if isinstance(result, dict): + # HO returns CreateImageDirectoryResponse{id}; the id may be at the + # top level of the operation-wait result or nested under "result". + dir_id = result.get("id") or (result.get("result") or {}).get("id") + if not dir_id: + raise CuttlefishError(f"create image dir: no 'id' in response: {result!r}") + return str(dir_id) + + def _populate_image_dir(self, dir_id: str, checksum: str) -> dict | list | str | None: + """Fill an image directory from a previously-uploaded artifact (async operation).""" + return self._do_operation("PUT", f"/cvd_imgs_dirs/{dir_id}", {"user_artifact_checksum": checksum}) + + def _stage_artifact(self, checksum: str, roles: set[str]) -> ImageDirGeneration: + """Populate a new staged generation and carry forward missing active roles.""" + generation = self._staged + if generation is None: + generation = ImageDirGeneration(dir_id=self._create_image_dir()) + self._staged = generation + + self._populate_image_dir(generation.dir_id, checksum) + for role in roles: + generation.artifacts[role] = checksum + + if self._active is not None: + for role, active_checksum in self._active.artifacts.items(): + if role not in generation.artifacts: + self._populate_image_dir(generation.dir_id, active_checksum) + generation.artifacts[role] = active_checksum + + return generation + + def _promote_staged(self) -> None: + """Make the staged generation active after its CVD is created.""" + staged = self._staged + if staged is None: + return + + previous = self._active + self._active = staged + self._staged = None + if previous is not None: + self._delete_image_dir_quietly(previous.dir_id) + + def _release_active(self) -> None: + """Delete the active generation after its CVD has been destroyed.""" + active = self._active + if active is not None and self._delete_image_dir_quietly(active.dir_id): + self._active = None + + def _env_config_for_create(self) -> dict: + """Build the env_config for POST /cvds, injecting the flashed image dir. + + The dir to boot from is the freshly-staged one if a flash is pending, + otherwise the dir the previous CVD used — so ``power off --destroy; + power on`` recreates the same flashed images without re-flashing. With + neither set, the configured env_config is used verbatim (e.g. + fetch-from-branch). + + HO's populate step symlinks each artifact's contents into a single image + dir (a per-filename merge), so one dir holds both the device images and + the host package even when they arrive as separate ``flash()`` calls. + For a complete bundle both path-valued fields point at the same + ``@image_dirs/{id}`` token — ``common.host_package`` (host tools) and + every ``instances[].disk.default_build`` (device images). A targeted + partial flash carries forward the active generation's checksums into + the new directory, so both fields continue to use the flashed + generation. HO rewrites the token to the on-host path + (``strings.ReplaceAll(config, "@image_dirs/", ...)``) before running + ``cvd load``. + """ + config = copy.deepcopy(self.env_config) + generation = self._staged or self._active + if generation is None: + return config + + token = f"@image_dirs/{generation.dir_id}" + if "host_package" in generation.artifacts: + config.setdefault("common", {})["host_package"] = token + if "images" in generation.artifacts: + instances = config.setdefault("instances", []) + if not instances: + instances.append({}) + for inst in instances: + inst.setdefault("disk", {})["default_build"] = token + return config + def _get_existing_cvds(self) -> list[dict]: """Return CVDs belonging to this driver's group. @@ -357,74 +543,87 @@ class CvdPower(VirtualPowerInterface, Driver): def client(cls) -> str: return "jumpstarter_driver_cuttlefish.client.CvdPowerClient" + def _delete_stale_cvds(self, cvds: list[dict]) -> None: + group_name = self.parent._cvd_group or self.parent.group + self.logger.warning("Found %d stale CVDs in group %s, deleting", len(cvds), group_name) + failed = [] + for cvd in cvds: + group = cvd.get("group", self.parent.group) + name = cvd.get("name", self.parent.name) + try: + self.parent._do_operation("DELETE", f"/cvds/{group}/{name}") + except CuttlefishError: + self.logger.warning("Failed to delete stale CVD %s/%s", group, name) + failed.append(f"{group}/{name}") + if failed: + raise CuttlefishError( + f"cannot create CVD - failed to delete stale CVDs: {', '.join(failed)}. " + f"Run 'j cuttlefish reset' then retry." + ) + + def _use_existing_cvd(self, cvd: dict) -> None: + self.parent._cvd_group = cvd.get("group") + self.parent._cvd_name = cvd.get("name") + self.logger.info( + "Found existing CVD %s/%s (status: %s)", + self.parent._cvd_group, + self.parent._cvd_name, + cvd.get("status"), + ) + if cvd.get("status") != "Running": + self.parent._do_operation("POST", f"{self.parent._cvd_path}/:start") + + def _create_cvd(self) -> None: + self.logger.info("Creating CVD from env_config") + try: + result = self.parent._do_operation( + "POST", + "/cvds", + {"env_config": self.parent._env_config_for_create()}, + timeout=600, + ) + except CuttlefishError as e: + msg = str(e) + if "in use" in msg or "already running" in msg or "ValidateTapDevices" in msg: + raise CuttlefishError( + f"CVD creation failed - orphaned processes from a previous session. " + f"Run 'j cuttlefish reset' then retry. Original error: {msg}" + ) from e + raise + + if isinstance(result, dict): + for cvd in result.get("cvds", []): + self.parent._cvd_group = cvd.get("group") + self.parent._cvd_name = cvd.get("name") + actual_port = cvd.get("adb_port") + if actual_port and actual_port != self.parent._expected_adb_port: + try: + self.parent._do_operation("DELETE", self.parent._cvd_path) + except CuttlefishError: + self.logger.warning("Failed to clean up CVD after port mismatch") + self.parent._cvd_group = None + self.parent._cvd_name = None + raise CuttlefishError( + f"HO assigned adb_port {actual_port} but expected " + f"{self.parent._expected_adb_port} — stale state may have leaked. " + f"Run 'j cuttlefish reset' then retry." + ) + break + + self.parent._promote_staged() + @export - def on(self) -> None: # noqa: C901 + def on(self) -> None: existing = self.parent._get_existing_cvds() if len(existing) > 1: - self.logger.warning( - "Found %d stale CVDs in group %s, deleting", len(existing), self.parent._cvd_group or self.parent.group - ) - failed = [] - for cvd in existing: - group = cvd.get("group", self.parent.group) - name = cvd.get("name", self.parent.name) - try: - self.parent._do_operation("DELETE", f"/cvds/{group}/{name}") - except CuttlefishError: - self.logger.warning("Failed to delete stale CVD %s/%s", group, name) - failed.append(f"{group}/{name}") - if failed: - raise CuttlefishError( - f"cannot create CVD - failed to delete stale CVDs: {', '.join(failed)}. " - f"Run 'j cuttlefish reset' then retry." - ) + self._delete_stale_cvds(existing) existing = [] if existing: - cvd = existing[0] - self.parent._cvd_group = cvd.get("group") - self.parent._cvd_name = cvd.get("name") - self.logger.info( - "Found existing CVD %s/%s (status: %s)", - self.parent._cvd_group, - self.parent._cvd_name, - cvd.get("status"), - ) - if cvd.get("status") != "Running": - self.parent._do_operation("POST", f"{self.parent._cvd_path}/:start") + self._use_existing_cvd(existing[0]) else: - self.logger.info("Creating CVD from env_config") - try: - result = self.parent._do_operation( - "POST", "/cvds", {"env_config": self.parent.env_config}, timeout=600, - ) - except CuttlefishError as e: - msg = str(e) - if "in use" in msg or "already running" in msg or "ValidateTapDevices" in msg: - raise CuttlefishError( - f"CVD creation failed - orphaned processes from a previous session. " - f"Run 'j cuttlefish reset' then retry. Original error: {msg}" - ) from e - raise - if isinstance(result, dict): - for cvd in result.get("cvds", []): - self.parent._cvd_group = cvd.get("group") - self.parent._cvd_name = cvd.get("name") - actual_port = cvd.get("adb_port") - if actual_port and actual_port != self.parent._expected_adb_port: - try: - self.parent._do_operation("DELETE", self.parent._cvd_path) - except CuttlefishError: - self.logger.warning("Failed to clean up CVD after port mismatch") - self.parent._cvd_group = None - self.parent._cvd_name = None - raise CuttlefishError( - f"HO assigned adb_port {actual_port} but expected " - f"{self.parent._expected_adb_port} — stale state may have leaked. " - f"Run 'j cuttlefish reset' then retry." - ) - break + self._create_cvd() self.parent._auto_connect_adb() if self.parent.boot_timeout: @@ -440,6 +639,7 @@ def off(self, destroy: bool = False) -> None: p._do_operation("DELETE", p._cvd_path) p._cvd_group = None p._cvd_name = None + p._release_active() else: self.logger.info(f"Stopping CVD {cvd_id}") p._do_operation("POST", f"{p._cvd_path}/:stop") @@ -451,16 +651,219 @@ def read(self) -> Generator[PowerReading, None, None]: @dataclass(kw_only=True) class CvdFlasher(FlasherInterface, Driver): - """Flasher for Cuttlefish devices (not yet implemented). - - Planned: upload artifacts to Host Orchestrator via its upload API. + """Flasher for Cuttlefish devices, backed by the HO user-artifact store. + + Each ``flash(source, target)`` streams one archive to the exporter, hashes + it (SHA-256), uploads it to HO's user-artifact store, extracts it, and + symlinks its contents into an image directory. HO's populate is a + *per-filename merge*, so several artifacts accumulate into one dir — the + dict form of the client API turns into one ``flash()`` call per entry, all + landing in the same staged dir:: + + client.storage.flash({ + "images": "aosp_cf_x86_64_auto-img.zip", # device images + "host_package": "cvd-host_package.tar.gz", # host tools + }) + + ``target`` must be either ``images`` (a ``.zip``) or ``host_package`` (a + ``.tar.gz``). An untargeted flash is treated as a complete bundle and + supplies both paths. The archive format is sniffed from magic bytes, since + HO's extractor dispatches on the upload's suffix and the resource stream + carries no name. + + A targeted flash may supply only one role; when an active generation exists, + the missing role is carried forward by checksum into the new immutable + generation. For a fresh CVD, provide both targeted artifacts (or one + complete untargeted bundle). + + Flash only STAGES — images are selected at CVD *creation*, so a plain + stop/start keeps the old images. The staged dir is consumed by the next + ``power on`` that creates a CVD:: + + j storage flash + j power off --destroy # stop/start would keep the old images + j power on # creates the CVD from the staged images """ parent: Cuttlefish + _VALID_TARGETS = {"images", "host_package"} + + def _validate_target(self, target: str | None) -> None: + if target is not None and target not in self._VALID_TARGETS: + raise CuttlefishError( + f"unsupported Cuttlefish storage target {target!r}; expected 'images' or 'host_package'" + ) + + def _roles_for_flash(self, target: str | None, suffix: str) -> set[str]: + """Validate a target label and return the roles this artifact supplies.""" + if target is None: + # An untargeted flash is the complete-bundle form and supplies both + # paths. Targeted calls are available for a split images/package + # upload or for preserving one path from env_config. + return set(self._VALID_TARGETS) + expected_suffix = ".zip" if target == "images" else ".tar.gz" + if suffix != expected_suffix: + raise CuttlefishError(f"storage target {target!r} requires {expected_suffix}, got {suffix}") + return {target} @export - def flash(self, source, target: str | None = None) -> None: - raise NotImplementedError("CvdFlasher.flash() not yet implemented") + async def flash(self, source, target: str | None = None) -> None: + self._validate_target(target) + tmp_path, checksum, size, suffix = await self._spool(source) + try: + roles = self._roles_for_flash(target, suffix) + label = target or "bundle" + if await to_thread.run_sync(self.parent._artifact_exists, checksum): + self.logger.info("artifact %s (%s) already present, skipping upload", checksum[:12], label) + else: + self.logger.info("uploading %s artifact %s (%d bytes)", label, checksum[:12], size) + await self._upload(checksum, tmp_path, size, suffix) + + self.logger.info("extracting artifact %s", checksum[:12]) + await to_thread.run_sync(self.parent._extract_artifact, checksum) + + had_staged = self.parent._staged is not None + generation = await to_thread.run_sync(self.parent._stage_artifact, checksum, roles) + if not had_staged: + self.logger.info("created image dir %s", generation.dir_id) + finally: + await to_thread.run_sync(os.unlink, tmp_path) + + self.logger.info( + "staged %s into image dir %s; create the CVD (power off --destroy; power on) to boot it", + label, + generation.dir_id, + ) + + def close(self): + """Reclaim pending staged image dirs at teardown. + + The active image dir is intentionally retained: the exporter can be + stopped while its CVD is still running, and deleting a directory that + backs that CVD can corrupt Host Orchestrator state. It is released by + ``power off --destroy`` after the CVD is deleted. HO has no delete route + for the underlying user artifacts, so those persist server-side by + design. + """ + p = self.parent + if p._staged: + p._delete_image_dir_quietly(p._staged.dir_id) + if p._active: + self.logger.warning( + "leaving active image dir %s at exporter teardown; destroy the CVD before manual cleanup", + p._active.dir_id, + ) + p._staged = None + p._active = None + super().close() + + # HO's extractor accepts only these two formats; both are detectable by the + # archive's leading magic bytes. + _GZIP_MAGIC = b"\x1f\x8b" + _ZIP_MAGIC = (b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08") # normal, empty, spanned + + async def _spool(self, source) -> tuple[str, str, int, str]: + """Stream `source` to a temp file, returning (path, sha256_hex, size, suffix). + + HO addresses uploads by checksum, so the full digest must be known + before the upload exists — hence spool-to-disk-then-upload rather than a + single streaming pass. The archive format is sniffed from the leading + magic bytes (the resource stream has no filename) so the upload can carry + the suffix HO's extractor dispatches on. + """ + fd, tmp_path = tempfile.mkstemp(prefix="cvd-flash-") + os.close(fd) + hasher = hashlib.sha256() + size = 0 + head = b"" + try: + async with await FileWriteStream.from_path(tmp_path) as out: + async with self.resource(source) as res: + async for chunk in res: + if len(head) < 4: + head += bytes(chunk[: 4 - len(head)]) + hasher.update(chunk) + size += len(chunk) + await out.send(chunk) + except BaseException: + try: + os.unlink(tmp_path) + except OSError: + pass + raise + suffix = self._sniff_suffix(head, tmp_path) + return tmp_path, hasher.hexdigest(), size, suffix + + def _sniff_suffix(self, head: bytes, tmp_path: str) -> str: + """Map an archive's magic bytes to the suffix HO's extractor requires.""" + if head.startswith(self._GZIP_MAGIC): + return ".tar.gz" + if head.startswith(self._ZIP_MAGIC): + return ".zip" + # Unusable spool — remove it here since flash()'s finally never runs when + # _spool raises, and fail clearly (HO would reject it anyway). + try: + os.unlink(tmp_path) + except OSError: + pass + raise CuttlefishError( + f"unsupported archive format: expected a .zip (device images) or a " + f".tar.gz (host package); leading bytes were {head[:4]!r}" + ) + + async def _upload(self, checksum: str, path: str, size: int, suffix: str) -> None: + """Upload a spooled artifact via HO's chunked multipart PUT. + + Each PUT carries one multipart ``file`` part plus ``chunk_offset_bytes`` + and ``file_size_bytes`` form fields; HO reassembles by offset. Uses + httpx (async) so a multi-GB transfer never blocks the event loop. + + A failed chunk is retried at the same offset — since HO reassembles by + offset, re-PUTting an offset is idempotent, so per-chunk retry doubles + as resume and a mid-transfer blip doesn't abort several GB of upload. + """ + url = f"{self.parent._base_url}/v1/userartifacts/{checksum}" + chunk_size = self.parent.upload_chunk_size + # The stored filename's suffix is what HO's extractor dispatches on, so + # it must reflect the real archive format (see _sniff_suffix). + filename = f"{checksum}{suffix}" + timeout = httpx.Timeout(connect=10.0, read=60.0, write=None, pool=10.0) + async with httpx.AsyncClient(timeout=timeout) as client: + async with await FileReadStream.from_path(path) as f: + offset = 0 + while offset < size: + buf = bytearray() + while len(buf) < chunk_size: + try: + buf += await f.receive(chunk_size - len(buf)) + except EndOfStream: + break + if not buf: + break + await self._put_chunk(client, url, filename, bytes(buf), offset, size) + offset += len(buf) + + async def _put_chunk( + self, client: "httpx.AsyncClient", url: str, filename: str, chunk: bytes, offset: int, size: int + ) -> None: + """PUT a single chunk, retrying at the same offset on transient failure.""" + attempts = max(1, self.parent.upload_max_retries) + for attempt in range(1, attempts + 1): + try: + resp = await client.put( + url, + files={"file": (filename, chunk, "application/octet-stream")}, + data={"chunk_offset_bytes": str(offset), "file_size_bytes": str(size)}, + ) + resp.raise_for_status() + return + except httpx.HTTPError as e: + if attempt >= attempts: + raise CuttlefishError(f"upload failed at offset {offset} after {attempts} attempts: {e}") from e + self.logger.warning( + "upload chunk at offset %d failed (attempt %d/%d): %s; retrying", offset, attempt, attempts, e + ) + await sleep(2 * attempt) @export def dump(self, target, partition: str | None = None) -> None: diff --git a/python/packages/jumpstarter-driver-cuttlefish/jumpstarter_driver_cuttlefish/driver_test.py b/python/packages/jumpstarter-driver-cuttlefish/jumpstarter_driver_cuttlefish/driver_test.py index 50207d0dd..7e9c5ab9e 100644 --- a/python/packages/jumpstarter-driver-cuttlefish/jumpstarter_driver_cuttlefish/driver_test.py +++ b/python/packages/jumpstarter-driver-cuttlefish/jumpstarter_driver_cuttlefish/driver_test.py @@ -1,11 +1,13 @@ import json +import logging +import os import subprocess -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest import requests -from .driver import Cuttlefish, CuttlefishError, CuttlefishTimeout +from .driver import Cuttlefish, CuttlefishError, CuttlefishTimeout, ImageDirGeneration BASE = "http://localhost:2080" @@ -139,6 +141,13 @@ def test_wait_unexpected_http_error(requests_mock, drv): drv._wait_for_operation("op-1") +def test_operation_wait_accepts_empty_success(requests_mock, drv): + requests_mock.delete(f"{BASE}/cvd_imgs_dirs/dir-1", json={"name": "op-1", "done": False}) + requests_mock.post(f"{BASE}/operations/op-1/:wait", status_code=200, text="") + + assert drv._do_operation("DELETE", "/cvd_imgs_dirs/dir-1") is None + + def _mock_op(requests_mock, method, path, op_name="op-1"): """Register mocks for an operation endpoint and its wait endpoint.""" getattr(requests_mock, method)(f"{BASE}{path}", json={"name": op_name, "done": False}) @@ -380,6 +389,34 @@ def test_cvd_power_off_destroy(requests_mock, drv): assert drv._cvd_name is None +def test_cvd_power_off_destroy_releases_active_generation(requests_mock, drv): + drv.children["adb"] = MagicMock() + drv._active = ImageDirGeneration("dir-active", {"images": "images-cs"}) + power = drv.children["power"] + requests_mock.delete(f"{BASE}/cvds/cvd_1/dev1", json={"name": "op-cvd", "done": False}) + requests_mock.post(f"{BASE}/operations/op-cvd/:wait", json={"name": "op-cvd", "done": True}) + requests_mock.delete(f"{BASE}/cvd_imgs_dirs/dir-active", json={"name": "op-dir", "done": False}) + requests_mock.post(f"{BASE}/operations/op-dir/:wait", json={"name": "op-dir", "done": True}) + + power.off(destroy=True) + + assert drv._active is None + assert any(r.method == "DELETE" and r.path == "/cvd_imgs_dirs/dir-active" for r in requests_mock.request_history) + + +def test_cvd_power_off_destroy_retains_active_generation_on_cleanup_failure(requests_mock, drv): + drv.children["adb"] = MagicMock() + drv._active = ImageDirGeneration("dir-active", {"images": "images-cs"}) + power = drv.children["power"] + requests_mock.delete(f"{BASE}/cvds/cvd_1/dev1", json={"name": "op-cvd", "done": False}) + requests_mock.post(f"{BASE}/operations/op-cvd/:wait", json={"name": "op-cvd", "done": True}) + requests_mock.delete(f"{BASE}/cvd_imgs_dirs/dir-active", status_code=500, json={"error": "temporary"}) + + power.off(destroy=True) + + assert drv._active == ImageDirGeneration("dir-active", {"images": "images-cs"}) + + def test_cvd_power_on_existing_running(requests_mock, drv): drv.children["adb"] = MagicMock() drv.boot_timeout = 0 @@ -487,16 +524,516 @@ def test_cvd_power_read_not_implemented(drv): list(power.read()) -def test_cvd_flasher_flash_not_implemented(drv): +def test_cvd_flasher_dump_not_implemented(drv): flasher = drv.children["storage"] with pytest.raises(NotImplementedError): - flasher.flash("source") + flasher.dump("target") -def test_cvd_flasher_dump_not_implemented(drv): +def test_env_config_for_create_passthrough(drv): + """With no staged image dirs, the configured env_config is returned as a copy.""" + drv.env_config = {"instances": [{"disk": {"default_build": "/home/vsoc-01/fetch"}}]} + result = drv._env_config_for_create() + assert result == drv.env_config + assert result is not drv.env_config # deep-copied, caller can't mutate config + result["instances"][0]["disk"]["default_build"] = "mutated" + assert drv.env_config["instances"][0]["disk"]["default_build"] == "/home/vsoc-01/fetch" + + +def test_env_config_for_create_injects_staged_dir(drv): + """A staged dir is referenced by both host_package and every default_build.""" + drv.env_config = {"instances": [{"disk": {}}, {"disk": {}}], "common": {}} + drv._staged = ImageDirGeneration("dir-7", {"images": "images-cs", "host_package": "host-cs"}) + result = drv._env_config_for_create() + assert result["common"]["host_package"] == "@image_dirs/dir-7" + assert all(inst["disk"]["default_build"] == "@image_dirs/dir-7" for inst in result["instances"]) + # source config is not mutated + assert drv.env_config["common"] == {} + + +def test_env_config_for_create_synthesizes_instance_when_empty(drv): + """With no configured instances, one is created so images are referenced.""" + drv.env_config = {} + drv._staged = ImageDirGeneration("dir-9", {"images": "images-cs", "host_package": "host-cs"}) + result = drv._env_config_for_create() + assert result["common"]["host_package"] == "@image_dirs/dir-9" + assert result["instances"] == [{"disk": {"default_build": "@image_dirs/dir-9"}}] + + +def test_env_config_for_create_falls_back_to_active_dir(drv): + """With nothing staged, the active dir is reused so recreate needs no reflash.""" + drv.env_config = {} + drv._active = ImageDirGeneration("dir-active", {"images": "images-cs", "host_package": "host-cs"}) + result = drv._env_config_for_create() + assert result["common"]["host_package"] == "@image_dirs/dir-active" + + +def test_env_config_for_create_prefers_staged_over_active(drv): + """A pending flash (staged) wins over the previous generation's active dir.""" + drv.env_config = {} + drv._active = ImageDirGeneration("dir-old", {"images": "images-cs", "host_package": "host-cs"}) + drv._staged = ImageDirGeneration("dir-new", {"images": "new-images-cs", "host_package": "new-host-cs"}) + result = drv._env_config_for_create() + assert result["common"]["host_package"] == "@image_dirs/dir-new" + + +class _FakeResource: + """Stand-in for Driver.resource(): async CM yielding an async byte iterator.""" + + def __init__(self, chunks): + self._chunks = chunks + + async def __aenter__(self): + async def gen(): + for c in self._chunks: + yield c + + return gen() + + async def __aexit__(self, *exc): + return False + + +class _FailingResource: + """Stand-in for a resource that fails while yielding its content.""" + + async def __aenter__(self): + async def gen(): + yield b"partial archive" + raise RuntimeError("stream failed") + + return gen() + + async def __aexit__(self, *exc): + return False + + +@pytest.mark.asyncio +async def test_flash_stages_image_dir(requests_mock, drv, monkeypatch): + """Happy path: existing artifact → extract → create + populate image dir → staged.""" + import hashlib + + payload = b"PK\x03\x04fake-image-zip" + checksum = hashlib.sha256(payload).hexdigest() flasher = drv.children["storage"] - with pytest.raises(NotImplementedError): - flasher.dump("target") + monkeypatch.setattr(flasher, "resource", lambda source: _FakeResource([payload])) + + # Artifact already present -> upload skipped (no httpx needed). + requests_mock.get(f"{BASE}/v1/userartifacts/{checksum}", status_code=200, json={}) + requests_mock.post(f"{BASE}/v1/userartifacts/{checksum}/:extract", json={"name": "op-x", "done": False}) + requests_mock.post(f"{BASE}/operations/op-x/:wait", json={"name": "op-x", "done": True}) + requests_mock.post(f"{BASE}/cvd_imgs_dirs", json={"name": "op-c", "done": False}) + requests_mock.post(f"{BASE}/operations/op-c/:wait", json={"name": "op-c", "done": True, "id": "dir-7"}) + requests_mock.put(f"{BASE}/cvd_imgs_dirs/dir-7", json={"name": "op-p", "done": False}) + requests_mock.post(f"{BASE}/operations/op-p/:wait", json={"name": "op-p", "done": True}) + + await flasher.flash("handle") + + # Staged, not yet active — the live dir only flips over on power.on(). + assert drv._staged == ImageDirGeneration("dir-7", {"images": checksum, "host_package": checksum}) + assert drv._active is None + + +@pytest.mark.asyncio +async def test_flash_uploads_when_artifact_absent(requests_mock, drv, monkeypatch): + """A 404 stat triggers the upload path with the streamed content's SHA-256.""" + import hashlib + + payload = b"PK\x03\x04another-image" + checksum = hashlib.sha256(payload).hexdigest() + flasher = drv.children["storage"] + monkeypatch.setattr(flasher, "resource", lambda source: _FakeResource([payload])) + + uploaded = {} + + async def fake_upload(cs, path, size, suffix): + uploaded["checksum"] = cs + uploaded["size"] = size + uploaded["suffix"] = suffix + + monkeypatch.setattr(flasher, "_upload", fake_upload) + + requests_mock.get(f"{BASE}/v1/userartifacts/{checksum}", status_code=404, json={"error": "not found"}) + requests_mock.post(f"{BASE}/v1/userartifacts/{checksum}/:extract", json={"name": "op-x", "done": False}) + requests_mock.post(f"{BASE}/operations/op-x/:wait", json={"name": "op-x", "done": True}) + requests_mock.post(f"{BASE}/cvd_imgs_dirs", json={"name": "op-c", "done": False}) + requests_mock.post(f"{BASE}/operations/op-c/:wait", json={"name": "op-c", "done": True, "id": "dir-9"}) + requests_mock.put(f"{BASE}/cvd_imgs_dirs/dir-9", json={"name": "op-p", "done": False}) + requests_mock.post(f"{BASE}/operations/op-p/:wait", json={"name": "op-p", "done": True}) + + await flasher.flash("handle") + + # The zip magic drives the stored suffix HO's extractor dispatches on. + assert uploaded == {"checksum": checksum, "size": len(payload), "suffix": ".zip"} + assert drv._staged == ImageDirGeneration("dir-9", {"images": checksum, "host_package": checksum}) + + +def test_create_image_dir_missing_id(requests_mock, drv): + """A create response with no id is a hard error, not a silent None.""" + requests_mock.post(f"{BASE}/cvd_imgs_dirs", json={"name": "op-c", "done": False}) + requests_mock.post(f"{BASE}/operations/op-c/:wait", json={"name": "op-c", "done": True}) + with pytest.raises(CuttlefishError, match="no 'id'"): + drv._create_image_dir() + + +def _mock_flash_endpoints(requests_mock, checksum, dir_id): + """Register the HO calls a single flash() makes (artifact already present).""" + requests_mock.get(f"{BASE}/v1/userartifacts/{checksum}", status_code=200, json={}) + requests_mock.post(f"{BASE}/v1/userartifacts/{checksum}/:extract", json={"name": "op-x", "done": False}) + requests_mock.post(f"{BASE}/operations/op-x/:wait", json={"name": "op-x", "done": True}) + requests_mock.post(f"{BASE}/cvd_imgs_dirs", json={"name": "op-c", "done": False}) + requests_mock.post(f"{BASE}/operations/op-c/:wait", json={"name": "op-c", "done": True, "id": dir_id}) + requests_mock.put(f"{BASE}/cvd_imgs_dirs/{dir_id}", json={"name": "op-p", "done": False}) + requests_mock.post(f"{BASE}/operations/op-p/:wait", json={"name": "op-p", "done": True}) + + +@pytest.mark.asyncio +async def test_flash_then_power_on_creates_from_staged_images(requests_mock, drv, monkeypatch): + """The full sequence: flash stages a dir, power.on() builds the CVD from it. + + This is the hop that a happy-path flash test alone would miss — power.on() + calls _env_config_for_create() unconditionally, so a broken injection only + surfaces here. + """ + import hashlib + + payload = b"PK\x03\x04complete-bundle" + checksum = hashlib.sha256(payload).hexdigest() + drv.env_config = {"instances": [{"disk": {}}], "common": {}} + drv.boot_timeout = 0 + drv.children["adb"] = MagicMock() + flasher = drv.children["storage"] + power = drv.children["power"] + monkeypatch.setattr(flasher, "resource", lambda source: _FakeResource([payload])) + + _mock_flash_endpoints(requests_mock, checksum, "dir-7") + await flasher.flash("handle") + + requests_mock.get(f"{BASE}/cvds", json={"cvds": []}) + requests_mock.post(f"{BASE}/cvds", json={"name": "op-cr", "done": False}) + requests_mock.post( + f"{BASE}/operations/op-cr/:wait", + json={"name": "op-cr", "done": True, "cvds": [{"group": "cvd_1", "name": "dev1", "adb_port": 6520}]}, + ) + power.on() + + create = [r for r in requests_mock.request_history if r.method == "POST" and r.path == "/cvds"][-1] + env = create.json()["env_config"] + assert env["common"]["host_package"] == "@image_dirs/dir-7" + assert env["instances"][0]["disk"]["default_build"] == "@image_dirs/dir-7" + + # Creating the CVD consumes the staged dir: it becomes the active (live) dir. + assert drv._staged is None + assert drv._active == ImageDirGeneration("dir-7", {"images": checksum, "host_package": checksum}) + + +@pytest.mark.asyncio +async def test_flash_two_artifacts_merge_into_one_dir(requests_mock, drv, monkeypatch): + """Two flash() calls before a create reuse one staged dir (HO merges on populate). + + This is the images + host_package split: `client.storage.flash({...})` + dispatches one flash() per entry, each PUT-populating the *same* dir with a + different checksum. Only the first call creates the dir. + """ + import hashlib + + images = b"PK\x03\x04device-images" + host_pkg = b"\x1f\x8bhost-package-tarball" + images_cs = hashlib.sha256(images).hexdigest() + host_cs = hashlib.sha256(host_pkg).hexdigest() + flasher = drv.children["storage"] + + payloads = iter([images, host_pkg]) + monkeypatch.setattr(flasher, "resource", lambda source: _FakeResource([next(payloads)])) + + # Only one dir is created (op-c → dir-1); each artifact PUT-populates it. + _mock_flash_endpoints(requests_mock, images_cs, "dir-1") + requests_mock.get(f"{BASE}/v1/userartifacts/{host_cs}", status_code=200, json={}) + requests_mock.post(f"{BASE}/v1/userartifacts/{host_cs}/:extract", json={"name": "op-x2", "done": False}) + requests_mock.post(f"{BASE}/operations/op-x2/:wait", json={"name": "op-x2", "done": True}) + + await flasher.flash("images-handle", target="images") + await flasher.flash("host-pkg-handle", target="host_package") + + assert drv._staged == ImageDirGeneration("dir-1", {"images": images_cs, "host_package": host_cs}) + # Exactly one dir created, but both artifacts populated into it. + creates = [r for r in requests_mock.request_history if r.method == "POST" and r.path == "/cvd_imgs_dirs"] + assert len(creates) == 1 + populates = [r for r in requests_mock.request_history if r.method == "PUT" and r.path == "/cvd_imgs_dirs/dir-1"] + assert len(populates) == 2 + assert {p.json()["user_artifact_checksum"] for p in populates} == {images_cs, host_cs} + assert drv._staged.artifacts == {"images": images_cs, "host_package": host_cs} + + +@pytest.mark.asyncio +async def test_partial_flash_carries_forward_active_artifacts(requests_mock, drv, monkeypatch): + """A partial flash creates an immutable generation with active artifacts copied by checksum.""" + import hashlib + + host_pkg = b"\x1f\x8bnew-host-package" + host_cs = hashlib.sha256(host_pkg).hexdigest() + old_images_cs = "old-images-checksum" + old_host_cs = "old-host-checksum" + drv._active = ImageDirGeneration("dir-old", {"images": old_images_cs, "host_package": old_host_cs}) + flasher = drv.children["storage"] + monkeypatch.setattr(flasher, "resource", lambda source: _FakeResource([host_pkg])) + + requests_mock.get(f"{BASE}/v1/userartifacts/{host_cs}", status_code=200, json={}) + requests_mock.post(f"{BASE}/v1/userartifacts/{host_cs}/:extract", json={"name": "op-x", "done": False}) + requests_mock.post(f"{BASE}/operations/op-x/:wait", json={"name": "op-x", "done": True}) + requests_mock.post(f"{BASE}/cvd_imgs_dirs", json={"name": "op-c", "done": False}) + requests_mock.post(f"{BASE}/operations/op-c/:wait", json={"name": "op-c", "done": True, "id": "dir-new"}) + requests_mock.put(f"{BASE}/cvd_imgs_dirs/dir-new", json={"name": "op-p", "done": False}) + requests_mock.post(f"{BASE}/operations/op-p/:wait", json={"name": "op-p", "done": True}) + + await flasher.flash("host-package", target="host_package") + + populates = [r.json()["user_artifact_checksum"] for r in requests_mock.request_history if r.method == "PUT"] + assert populates == [host_cs, old_images_cs] + assert drv._staged == ImageDirGeneration("dir-new", {"host_package": host_cs, "images": old_images_cs}) + env = drv._env_config_for_create() + assert env["common"]["host_package"] == "@image_dirs/dir-new" + assert env["instances"][0]["disk"]["default_build"] == "@image_dirs/dir-new" + + +@pytest.mark.asyncio +async def test_put_chunk_retries_same_offset(drv): + """A failed chunk is retried with the same offset and payload.""" + import httpx + + class FakeResponse: + def raise_for_status(self): + return None + + class FakeClient: + def __init__(self): + self.calls = [] + + async def put(self, url, *, files, data): + self.calls.append((url, files, data)) + if len(self.calls) == 1: + raise httpx.ReadError("transient failure") + return FakeResponse() + + client = FakeClient() + flasher = drv.children["storage"] + with patch("jumpstarter_driver_cuttlefish.driver.sleep", new=AsyncMock()): + await flasher._put_chunk(client, "http://localhost/upload", "artifact.zip", b"chunk", 16, 100) + + assert len(client.calls) == 2 + assert ( + client.calls[0][2] + == client.calls[1][2] + == { + "chunk_offset_bytes": "16", + "file_size_bytes": "100", + } + ) + assert client.calls[0][1]["file"][1] == client.calls[1][1]["file"][1] == b"chunk" + + +@pytest.mark.asyncio +async def test_spool_removes_partial_file_on_stream_failure(drv, tmp_path, monkeypatch): + """A failed source stream does not leave a partial spool file behind.""" + spool_path = tmp_path / "cvd-flash-failed" + + def fake_mkstemp(prefix): + return os.open(spool_path, os.O_RDWR | os.O_CREAT | os.O_EXCL), str(spool_path) + + flasher = drv.children["storage"] + monkeypatch.setattr("jumpstarter_driver_cuttlefish.driver.tempfile.mkstemp", fake_mkstemp) + monkeypatch.setattr(flasher, "resource", lambda source: _FailingResource()) + + with pytest.raises(RuntimeError, match="stream failed"): + await flasher._spool("handle") + + assert not spool_path.exists() + + +@pytest.mark.asyncio +async def test_upload_uses_bounded_network_timeouts(drv, tmp_path, monkeypatch): + """Upload connections have finite control-plane timeouts and unbounded writes.""" + import httpx + + class FakeResponse: + def raise_for_status(self): + return None + + class FakeClient: + instances = [] + + def __init__(self, *, timeout): + self.timeout = timeout + self.instances.append(self) + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + async def put(self, url, *, files, data): + return FakeResponse() + + monkeypatch.setattr("jumpstarter_driver_cuttlefish.driver.httpx.AsyncClient", FakeClient) + path = tmp_path / "artifact.zip" + path.write_bytes(b"artifact") + + await drv.children["storage"]._upload("checksum", str(path), path.stat().st_size, ".zip") + + timeout = FakeClient.instances[0].timeout + assert isinstance(timeout, httpx.Timeout) + assert timeout.connect == 10.0 + assert timeout.read == 60.0 + assert timeout.write is None + assert timeout.pool == 10.0 + + +def test_targeted_generation_injects_only_recorded_roles(drv): + """A generation injects only the roles it contains.""" + drv.env_config = { + "instances": [{"disk": {"default_build": "/old/images"}}], + "common": {"host_package": "/old/host"}, + } + drv._staged = ImageDirGeneration("dir-images", {"images": "images-cs"}) + + result = drv._env_config_for_create() + + assert result["instances"][0]["disk"]["default_build"] == "@image_dirs/dir-images" + assert result["common"]["host_package"] == "/old/host" + + +@pytest.mark.asyncio +async def test_flash_rejects_unknown_target(drv, monkeypatch): + """Cuttlefish storage exposes only the documented target roles.""" + flasher = drv.children["storage"] + resource = MagicMock() + monkeypatch.setattr(flasher, "resource", resource) + + with pytest.raises(CuttlefishError, match="expected 'images' or 'host_package'"): + await flasher.flash("handle", target="system") + resource.assert_not_called() + + +@pytest.mark.asyncio +async def test_flash_rejects_target_format_mismatch(drv, monkeypatch): + """The documented target role must match the archive format.""" + flasher = drv.children["storage"] + monkeypatch.setattr(flasher, "resource", lambda source: _FakeResource([b"\x1f\x8bhost-package"])) + + with pytest.raises(CuttlefishError, match=r"requires \.zip"): + await flasher.flash("handle", target="images") + + +@pytest.mark.asyncio +async def test_flash_after_create_starts_new_generation(requests_mock, drv, monkeypatch): + """A flash after a create begins a fresh staged dir, leaving the active one intact.""" + import hashlib + + payload = b"PK\x03\x04new-generation" + checksum = hashlib.sha256(payload).hexdigest() + flasher = drv.children["storage"] + monkeypatch.setattr(flasher, "resource", lambda source: _FakeResource([payload])) + # Simulate a prior generation already booted from dir-old. + drv._active = ImageDirGeneration("dir-old", {"images": "old-images", "host_package": "old-host"}) + + _mock_flash_endpoints(requests_mock, checksum, "dir-new") + await flasher.flash("handle") + + # New dir staged; the live dir is untouched until the next create consumes it. + assert drv._staged == ImageDirGeneration("dir-new", {"images": checksum, "host_package": checksum}) + assert drv._active == ImageDirGeneration("dir-old", {"images": "old-images", "host_package": "old-host"}) + assert not any(r.method == "DELETE" for r in requests_mock.request_history) + + +@pytest.mark.asyncio +async def test_extract_tolerates_already_extracted(requests_mock, drv, monkeypatch): + """A 409 from :extract (artifact already extracted) is treated as success.""" + import hashlib + + payload = b"PK\x03\x04already-extracted" + checksum = hashlib.sha256(payload).hexdigest() + flasher = drv.children["storage"] + monkeypatch.setattr(flasher, "resource", lambda source: _FakeResource([payload])) + + requests_mock.get(f"{BASE}/v1/userartifacts/{checksum}", status_code=200, json={}) + requests_mock.post(f"{BASE}/v1/userartifacts/{checksum}/:extract", status_code=409, json={"error": "extracted"}) + requests_mock.post(f"{BASE}/cvd_imgs_dirs", json={"name": "op-c", "done": False}) + requests_mock.post(f"{BASE}/operations/op-c/:wait", json={"name": "op-c", "done": True, "id": "dir-7"}) + requests_mock.put(f"{BASE}/cvd_imgs_dirs/dir-7", json={"name": "op-p", "done": False}) + requests_mock.post(f"{BASE}/operations/op-p/:wait", json={"name": "op-p", "done": True}) + + await flasher.flash("handle") + + assert drv._staged == ImageDirGeneration("dir-7", {"images": checksum, "host_package": checksum}) + + +def test_extract_tolerates_already_extracted_from_wait(requests_mock, drv): + checksum = "abc123" + requests_mock.post( + f"{BASE}/v1/userartifacts/{checksum}/:extract", + json={"name": "op-x", "done": False}, + ) + requests_mock.post( + f"{BASE}/operations/op-x/:wait", + status_code=409, + json={"error": "already extracted"}, + ) + + drv._extract_artifact(checksum) + + +@pytest.mark.asyncio +async def test_spool_sniffs_gzip_suffix(drv, monkeypatch): + """A gzip-magic stream is spooled with a .tar.gz suffix (host package).""" + flasher = drv.children["storage"] + monkeypatch.setattr(flasher, "resource", lambda source: _FakeResource([b"\x1f\x8bhost-package"])) + tmp_path, _checksum, _size, suffix = await flasher._spool("handle") + try: + assert suffix == ".tar.gz" + finally: + os.unlink(tmp_path) + + +@pytest.mark.asyncio +async def test_spool_rejects_unknown_format(drv, monkeypatch): + """A stream with no recognized archive magic is rejected and its spool removed.""" + flasher = drv.children["storage"] + monkeypatch.setattr(flasher, "resource", lambda source: _FakeResource([b"not-an-archive"])) + with pytest.raises(CuttlefishError, match="unsupported archive format"): + await flasher._spool("handle") + + +def test_flasher_close_reclaims_staged_but_retains_active_dir(requests_mock, drv, caplog): + """close() cleans pending work but retains a dir a live CVD may reference.""" + flasher = drv.children["storage"] + drv._staged = ImageDirGeneration("dir-staged", {"images": "images-cs"}) + drv._active = ImageDirGeneration("dir-active", {"images": "images-cs", "host_package": "host-cs"}) + requests_mock.delete(f"{BASE}/cvd_imgs_dirs/dir-staged", json={"name": "op-staged", "done": False}) + requests_mock.post(f"{BASE}/operations/op-staged/:wait", json={"name": "op-staged", "done": True}) + + caplog.set_level(logging.WARNING) + flasher.close() + + assert drv._staged is None + assert drv._active is None + deleted = {r.path for r in requests_mock.request_history if r.method == "DELETE"} + assert deleted == {"/cvd_imgs_dirs/dir-staged"} + assert "leaving active image dir dir-active at exporter teardown" in caplog.text + + +def test_flasher_close_survives_delete_failure(requests_mock, drv, caplog): + """A failed cleanup DELETE is logged, not raised, and state is still cleared.""" + flasher = drv.children["storage"] + drv._staged = ImageDirGeneration("dir-7", {"images": "images-cs"}) + requests_mock.delete(f"{BASE}/cvd_imgs_dirs/dir-7", status_code=500, json={"error": "in use"}) + + caplog.set_level(logging.WARNING) + flasher.close() + + assert drv._staged is None + assert drv._active is None + assert "failed to delete image dir dir-7; it may need manual cleanup" in caplog.text def test_cvd_power_on_ignores_other_groups(requests_mock, drv): diff --git a/python/packages/jumpstarter-driver-cuttlefish/pyproject.toml b/python/packages/jumpstarter-driver-cuttlefish/pyproject.toml index da3b5d3cc..004967d5c 100644 --- a/python/packages/jumpstarter-driver-cuttlefish/pyproject.toml +++ b/python/packages/jumpstarter-driver-cuttlefish/pyproject.toml @@ -6,7 +6,9 @@ license = "Apache-2.0" authors = [{ name = "Benny Zlotnik", email = "bzlotnik@redhat.com" }] requires-python = ">=3.11" dependencies = [ + "anyio>=4.0.0", "click>=8.0.0", + "httpx>=0.27.0", "jumpstarter", "jumpstarter-driver-adb", "jumpstarter-driver-composite",