Source code for tests.utils.fwu_utils

# SPDX-FileCopyrightText: <text>Copyright 2026 Arm Limited and/or its
# affiliates <open-source-office@arm.com></text>
#
# SPDX-License-Identifier: MIT

import pytest
import time

from tests.utils.session_utils import expect_pattern_live
from tests.utils.common_utils import CommonUtils



logger = CommonUtils.get_logger(__name__)





def check_block_device(platform_base_obj, dev: str) -> None:
    """
    Verify that a block device exists on target.

    :param platform_base_obj: Platform fixture containing manager
    :param dev: Device path
    :returns: None
    """
    cmd = f"lsblk -no NAME {dev}"
    mgr = platform_base_obj.mgr
    status, output = mgr.execute_command_with_prompt_capture(
        port=platform_base_obj.default_console,
        command=cmd,
        timeout=50,
    )

    assert status == 0, "Failed to call lsblk to list block devices"

    # Output is non-empty if the block device exists
    assert output.strip() != "", f"Block device {dev} missing on target"





def check_mount(
    platform_base_obj, target: str, source: str, error_msg: str
) -> None:
    """
    Verify that a mount exists using findmnt linux command

    :param platform_base_obj: Platform fixture containing manager
    :param target: target to mount the device
    :param source: Expected device
    :param error_msg: Msg in case of failure
    :returns: None
    """
    cmd = f"findmnt -T {target} -o target,source -n"
    mgr = platform_base_obj.mgr
    status, output = mgr.execute_command_with_prompt_capture(
        port=platform_base_obj.default_console,
        command=cmd,
        timeout=120,
    )

    # Non-zero exit means findmnt itself failed
    assert status == 0, f"findmnt failed for {target} with status {status}"

    if target not in output or source not in output:
        pytest.fail(error_msg)





def check_file(platform_base_obj, path: str, error_msg: str) -> None:
    """
    Verify that a file exists on target.

    :param platform_base_obj: Platform fixture containing manager
    :param path: File path
    :param error_msg: Msg in case of failure
    :returns: None
    """
    cmd = f"test -f {path}"
    mgr = platform_base_obj.mgr
    status, _ = mgr.execute_command_with_prompt_capture(
        port=platform_base_obj.default_console,
        command=cmd,
        timeout=60,
    )
    assert status == 0, error_msg





def prepare_capsule_environment(platform_base_obj):
    """
    Mount boot, verify devices, check capsule, and copy file
    """
    mgr = platform_base_obj.mgr
    # Verify that both block devices exist before mounting
    for dev in ("/dev/vda1", "/dev/vdb1"):
        check_block_device(platform_base_obj, dev)

    # Mount boot and capsule source partitions
    mounts = (
        ("/dev/vda1", "/boot"),
        ("/dev/vdb1", "/mnt"),
    )

    for device, mount_point in mounts:
        status, output = mgr.execute_command_with_prompt_capture(
            port=platform_base_obj.default_console,
            command=f"mount {device} {mount_point}",
            timeout=240,
        )
        assert status == 0, f"Failed to mount {device} to {mount_point}."

        check_mount(
            platform_base_obj,
            mount_point,
            device,
            f"Mount check failed: {device} not mounted on {mount_point}",
        )

    # Verify capsule presence before copy
    check_file(
        platform_base_obj,
        "/mnt/fw.cap",
        "Capsule /mnt/fw.cap missing before copy",
    )
    status, output = mgr.execute_command_with_prompt_capture(
        port=platform_base_obj.default_console,
        command=r"mkdir -p /boot/EFI/UpdateCapsule",
        timeout=60,
    )

    assert status == 0, "Failed to create /boot/EFI/UpdateCapsule"

    status, output = mgr.execute_command_with_prompt_capture(
        port=platform_base_obj.default_console,
        command="cp -rf /mnt/fw.cap /boot/EFI/UpdateCapsule/",
        timeout=60,
    )

    assert status == 0, "Failed to copy capsule to /EFI/UpdateCapsule"

    # Verify the copy succeeded
    check_file(
        platform_base_obj,
        "/boot/EFI/UpdateCapsule/fw.cap",
        "Capsule copy to EFI UpdateCapsule failed",
    )





def check_system_status(platform_base_obj, timeout: int = 200) -> None:
    """
    Verify that systemd reaches ``running`` state after FWU flow.

    If running state is not reached, this helper drains pending console output,
    collects failed systemd units, and fails the test with diagnostics.

    :param platform_base_obj: Platform fixture containing manager
    :param timeout: Reserved timeout argument
    :returns: None
    """
    mgr = platform_base_obj.mgr
    linux_console = platform_base_obj.default_console
    session = mgr.sessions[linux_console]

    # send command to check system status
    cmd = "systemctl is-system-running --wait"
    session.sendline(cmd)

    # Expect 'running' status
    status = expect_pattern_live(
        mgr.sessions[linux_console],
        r"running",
        timeout=240,
        msg="Systemd is not running after FWU",
        level="error",
    )
    time.sleep(2)

    """ If systemd did not reach 'running',clear the buffer.
    Then capture failed units to help debug the failure and log them.
    """
    if status:
        return
    else:
        mgr.drain_console_to_tail(session, timeout=100)
        time.sleep(1)
        session.sendline("systemctl list-units --failed")
        time.sleep(2)

    failed_output = session.before or ""
    msg = (
        "Systemd did not reach 'running' state.\n"
        f"Failed units:\n{failed_output}"
    )
    pytest.fail(msg)