NODEDC_MISSION_CORE/tests/test_usb_snapshot.py

from __future__ import annotations

import json
import plistlib
from pathlib import Path

import pytest
from typer.testing import CliRunner

from k1link.device_plugins.xgrids_k1.cli import app
from k1link.device_plugins.xgrids_k1.usb.snapshot import (
    DISKUTIL_COMMAND,
    SERIAL_IOREG_COMMAND,
    USB_IOREG_COMMAND,
    CommandOutput,
    snapshot,
)

runner = CliRunner()


def _plist_output(argv: tuple[str, ...], payload: object) -> CommandOutput:
    return CommandOutput(
        argv=argv,
        returncode=0,
        stdout=plistlib.dumps(payload),
        stderr=b"",
        error=None,
    )


def _interface(
    name: str,
    number: int,
    interface_class: int,
    subclass: int,
    protocol: int,
) -> dict[str, object]:
    return {
        "IOObjectClass": "IOUSBHostInterface",
        "IORegistryEntryName": name,
        "kUSBString": name,
        "bInterfaceNumber": number,
        "bInterfaceClass": interface_class,
        "bInterfaceSubClass": subclass,
        "bInterfaceProtocol": protocol,
        "bAlternateSetting": 0,
        "bConfigurationValue": 1,
        "bNumEndpoints": 2,
    }


def test_snapshot_parses_xgrids_interfaces_storage_and_usbmodem() -> None:
    usb_plist = [
        {
            "IOObjectClass": "IOUSBHostDevice",
            "IORegistryEntryName": "XGRIDS Device",
            "USB Product Name": "XGRIDS Device",
            "USB Vendor Name": "rockchip",
            "USB Serial Number": "synthetic-serial",
            "idVendor": 0x2207,
            "idProduct": 0x0019,
            "bDeviceClass": 0xEF,
            "bDeviceSubClass": 0x02,
            "bDeviceProtocol": 0x01,
            "bcdUSB": 0x0210,
            "bcdDevice": 0x0310,
            "USBSpeed": 3,
            "UsbLinkSpeed": 480_000_000,
            "USB Address": 1,
            "locationID": 0x01100000,
            "IORegistryEntryID": 12345,
            "IORegistryEntryChildren": [
                _interface("RNDIS Communications Control", 0, 0xE0, 0x01, 0x03),
                _interface("RNDIS Ethernet Data", 1, 0x0A, 0x00, 0x00),
                _interface("Mass Storage", 2, 0x08, 0x06, 0x50),
                _interface("CDC NCM Control", 3, 0x02, 0x0D, 0x00),
                _interface("CDC ACM Serial", 4, 0x02, 0x02, 0x01),
                {
                    "IOObjectClass": "IOMedia",
                    "IORegistryEntryName": "Synthetic Media",
                    "BSD Name": "disk4",
                },
            ],
        },
        {
            "IOObjectClass": "IOUSBHostDevice",
            "IORegistryEntryName": "Unrelated Camera",
            "USB Product Name": "Unrelated Camera",
            "idVendor": 999,
        },
    ]
    serial_plist = [
        {
            "IOCalloutDevice": "/dev/cu.usbmodemK1TEST",
            "IODialinDevice": "/dev/tty.usbmodemK1TEST",
            "IOTTYBaseName": "usbmodemK1TEST",
            "IOSerialBSDClientType": "IOSerialStream",
        },
        {
            "IOCalloutDevice": "/dev/cu.debug-console",
            "IODialinDevice": "/dev/tty.debug-console",
        },
    ]
    storage_plist = {
        "AllDisks": ["disk4", "disk4s1"],
        "WholeDisks": ["disk4"],
        "VolumesFromDisks": ["K1_DATA"],
        "AllDisksAndPartitions": [
            {
                "DeviceIdentifier": "disk4",
                "Content": "GUID_partition_scheme",
                "Size": 64_000_000,
                "OSInternal": False,
                "Partitions": [
                    {
                        "DeviceIdentifier": "disk4s1",
                        "Content": "Microsoft Basic Data",
                        "Size": 63_000_000,
                        "VolumeName": "K1_DATA",
                        "MountPoint": "/Volumes/K1_DATA",
                        "OSInternal": False,
                    }
                ],
            }
        ],
    }
    outputs = {
        USB_IOREG_COMMAND: _plist_output(USB_IOREG_COMMAND, usb_plist),
        SERIAL_IOREG_COMMAND: _plist_output(SERIAL_IOREG_COMMAND, serial_plist),
        DISKUTIL_COMMAND: _plist_output(DISKUTIL_COMMAND, storage_plist),
    }

    result = snapshot(lambda argv: outputs[tuple(argv)])

    assert result["xgrids_device_count"] == 1
    device = result["xgrids_devices"][0]
    assert device["vendor_id_hex"] == "0x2207"
    assert device["product_id_hex"] == "0x0019"
    assert device["bsd_names"] == ["disk4"]
    assert device["interface_capabilities"] == [
        "cdc_data",
        "mass_storage",
        "ncm",
        "rndis",
        "serial",
    ]
    assert [interface["interface_number"] for interface in device["interfaces"]] == [
        0,
        1,
        2,
        3,
        4,
    ]
    assert all(entry["xgrids_related"] for entry in result["external_storage"]["entries"])
    assert result["usbmodem_device_names"] == [
        "/dev/cu.usbmodemK1TEST",
        "/dev/tty.usbmodemK1TEST",
    ]
    assert len(result["usbmodem_devices"]) == 1
    assert all(source["ok"] for source in result["sources"])
    assert all("sudo" not in source["argv"] for source in result["sources"])
    assert result["safety"] == {
        "metadata_only": True,
        "sudo_used": False,
        "device_file_contents_read": False,
        "device_writes_performed": False,
    }


def test_snapshot_reports_command_and_plist_errors_without_raising() -> None:
    outputs = {
        USB_IOREG_COMMAND: CommandOutput(
            argv=USB_IOREG_COMMAND,
            returncode=0,
            stdout=b"not a plist",
            stderr=b"",
            error=None,
        ),
        SERIAL_IOREG_COMMAND: CommandOutput(
            argv=SERIAL_IOREG_COMMAND,
            returncode=1,
            stdout=b"",
            stderr=b"serial registry unavailable",
            error=None,
        ),
        DISKUTIL_COMMAND: _plist_output(
            DISKUTIL_COMMAND,
            {
                "AllDisks": [],
                "WholeDisks": [],
                "VolumesFromDisks": [],
                "AllDisksAndPartitions": [],
            },
        ),
    }

    result = snapshot(lambda argv: outputs[tuple(argv)])

    assert result["xgrids_devices"] == []
    statuses = {source["name"]: source for source in result["sources"]}
    assert statuses["usb_ioreg"]["ok"] is False
    assert statuses["usb_ioreg"]["error"].startswith("invalid plist:")
    assert statuses["serial_ioreg"]["ok"] is False
    assert statuses["serial_ioreg"]["error"] == "serial registry unavailable"
    assert statuses["external_disks"]["ok"] is True


def test_snapshot_treats_empty_ioreg_output_as_no_devices() -> None:
    outputs = {
        USB_IOREG_COMMAND: CommandOutput(USB_IOREG_COMMAND, 0, b"", b"", None),
        SERIAL_IOREG_COMMAND: CommandOutput(SERIAL_IOREG_COMMAND, 0, b"", b"", None),
        DISKUTIL_COMMAND: _plist_output(
            DISKUTIL_COMMAND,
            {
                "AllDisks": [],
                "WholeDisks": [],
                "VolumesFromDisks": [],
                "AllDisksAndPartitions": [],
            },
        ),
    }

    result = snapshot(lambda argv: outputs[tuple(argv)])

    assert result["xgrids_device_count"] == 0
    assert all(source["ok"] for source in result["sources"])


def test_usb_snapshot_cli_writes_json(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    output = tmp_path / "nested" / "usb-snapshot.json"
    payload = {
        "schema_version": 1,
        "xgrids_device_count": 0,
        "safety": {
            "metadata_only": True,
            "sudo_used": False,
            "device_file_contents_read": False,
            "device_writes_performed": False,
        },
    }
    monkeypatch.setattr("k1link.device_plugins.xgrids_k1.cli.usb_snapshot", lambda: payload)

    result = runner.invoke(app, ["usb", "snapshot", "--out", str(output)])

    assert result.exit_code == 0
    assert json.loads(output.read_text(encoding="utf-8")) == payload
    assert "no sudo" in result.stdout