jumpstarter/python/packages/jumpstarter-driver-uboot/jumpstarter_driver_uboot/client_test.py at 900c12af707f72eabdd514edc1a23d5068c40544 · jumpstarter-dev/jumpstarter · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
import inspect
import logging
from contextlib import contextmanager
from unittest.mock import MagicMock, call, patch

import pexpect
import pytest

from .client import UbootConsoleClient
from .common import ESC


def test_reboot_to_console_accepts_retries_kwarg() -> None:
    sig = inspect.signature(UbootConsoleClient.reboot_to_console)
    retries_param = sig.parameters.get("retries")
    assert retries_param is not None
    assert retries_param.default == 100
    assert retries_param.kind == inspect.Parameter.KEYWORD_ONLY


def test_reboot_to_console_retries_limits_attempts() -> None:
    mock_power = MagicMock()
    mock_pexpect_process = MagicMock()
    mock_pexpect_process.send = MagicMock()
    mock_pexpect_process.expect_exact = MagicMock(side_effect=pexpect.TIMEOUT("timeout"))

    @contextmanager
    def fake_pexpect():
        yield mock_pexpect_process

    mock_serial = MagicMock()
    mock_serial.pexpect = fake_pexpect

    client = object.__new__(UbootConsoleClient)
    client.children = {"power": mock_power, "serial": mock_serial}
    client.logger = logging.getLogger("test_uboot")

    prompt_value = "=> "
    with patch.object(type(client), "prompt", new_callable=lambda: property(lambda self: prompt_value)):
        with pytest.raises(RuntimeError, match="Failed to get U-Boot prompt"):
            with client.reboot_to_console(retries=3):
                pass

    assert mock_pexpect_process.send.call_count == 3
    mock_pexpect_process.send.assert_has_calls([call(ESC)] * 3)
    mock_power.cycle.assert_called_once()


def test_reboot_to_console_yields_on_prompt_match() -> None:
    mock_power = MagicMock()
    mock_pexpect_process = MagicMock()
    mock_pexpect_process.send = MagicMock()
    mock_pexpect_process.expect_exact = MagicMock(
        side_effect=[pexpect.TIMEOUT("timeout"), pexpect.TIMEOUT("timeout"), None]
    )

    @contextmanager
    def fake_pexpect():
        yield mock_pexpect_process

    mock_serial = MagicMock()
    mock_serial.pexpect = fake_pexpect

    client = object.__new__(UbootConsoleClient)
    client.children = {"power": mock_power, "serial": mock_serial}
    client.logger = logging.getLogger("test_uboot")

    prompt_value = "=> "
    with patch.object(type(client), "prompt", new_callable=lambda: property(lambda self: prompt_value)):
        entered = False
        with client.reboot_to_console(retries=5):
            entered = True

    assert entered
    assert mock_pexpect_process.send.call_count == 3
    mock_power.cycle.assert_called_once()


def test_reboot_to_console_retries_zero_raises_immediately() -> None:
    mock_power = MagicMock()
    mock_pexpect_process = MagicMock()

    @contextmanager
    def fake_pexpect():
        yield mock_pexpect_process

    mock_serial = MagicMock()
    mock_serial.pexpect = fake_pexpect

    client = object.__new__(UbootConsoleClient)
    client.children = {"power": mock_power, "serial": mock_serial}
    client.logger = logging.getLogger("test_uboot")

    prompt_value = "=> "
    with patch.object(type(client), "prompt", new_callable=lambda: property(lambda self: prompt_value)):
        with pytest.raises(RuntimeError, match="Failed to get U-Boot prompt"):
            with client.reboot_to_console(retries=0):
                pass

    mock_pexpect_process.send.assert_not_called()
    mock_power.cycle.assert_called_once()