diff --git a/src/ci_orchestrator/commands/template.py b/src/ci_orchestrator/commands/template.py index aabbe40..eeba90a 100644 --- a/src/ci_orchestrator/commands/template.py +++ b/src/ci_orchestrator/commands/template.py @@ -292,6 +292,44 @@ def _wait_tcp(host: str, port: int, timeout: float) -> bool: return False +def _wait_tcp_down(host: str, port: int, timeout: float, interval: float = 5.0) -> bool: + """Return True once TCP *port* on *host* stops accepting connections (host rebooting). + + Used after issuing a reboot: a pending Windows Update / CBS install can delay + the actual restart for minutes while the WinRM session is still alive, so we + must observe the port go *down* before waiting for it to come back — otherwise + we'd treat the still-up pre-reboot session as "ready". + """ + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + with socket.create_connection((host, port), timeout=5.0): + time.sleep(interval) + except OSError: + return True + return False + + +def _reboot_and_wait( + transport: WinRmTransport, host: str, port: int, winrm_timeout: int +) -> None: + """Reboot the guest and wait for it to actually go down then come back ready.""" + with contextlib.suppress(TransportError): + transport.run("Restart-Computer -Force", check=False) + transport.close() + click.echo("[prepare-win] waiting for VM to go down ...") + if _wait_tcp_down(host, port, timeout=300.0): + click.echo("[prepare-win] VM down — waiting for it to come back ...") + else: + click.echo( + "[prepare-win] [WARN] WinRM still reachable after reboot request; " + "proceeding anyway." + ) + if not _wait_tcp(host, port, float(winrm_timeout)): + raise click.ClickException(f"TCP/{port} on {host} not reachable after reboot.") + _wait_winrm(transport, winrm_timeout) + + def _parse_exit_marker(stdout: str) -> int: """Extract ``CI_EXITCODE:N`` written by the invoke wrapper; default 0.""" for line in reversed(stdout.splitlines()): @@ -582,12 +620,7 @@ def template_prepare_win( ) click.echo("[prepare-win] Windows Update requires reboot — rebooting ...") - with contextlib.suppress(TransportError): - transport.run("Restart-Computer -Force", check=False) - transport.close() - time.sleep(30) - click.echo(f"[prepare-win] Waiting for WinRM after reboot ({winrm_timeout}s) ...") - _wait_winrm(transport, winrm_timeout) + _reboot_and_wait(transport, guest_ip, 5986, winrm_timeout) click.echo("[prepare-win] WinRM ready after reboot.") # Install CI-StaticIp scheduled task @@ -653,12 +686,7 @@ Register-ScheduledTask -TaskName 'CI-StaticIp' ` # Final reboot to flush installation state click.echo("[prepare-win] Final reboot to flush installation state ...") - with contextlib.suppress(TransportError): - transport.run("Restart-Computer -Force", check=False) - transport.close() - time.sleep(30) - click.echo("[prepare-win] Waiting for WinRM after final reboot ...") - _wait_winrm(transport, winrm_timeout) + _reboot_and_wait(transport, guest_ip, 5986, winrm_timeout) # Disk cleanup — runs after final reboot on fully-installed system if not skip_cleanup: diff --git a/tests/python/test_commands_template.py b/tests/python/test_commands_template.py index bba2bb1..7e1495f 100644 --- a/tests/python/test_commands_template.py +++ b/tests/python/test_commands_template.py @@ -363,6 +363,47 @@ def test_wait_tcp_timeout(monkeypatch: pytest.MonkeyPatch) -> None: assert tmpl_mod._wait_tcp("127.0.0.1", 9999, 0.0) is False +def test_wait_tcp_down_detects_down(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + tmpl_mod.socket, "create_connection", + lambda *a, **kw: (_ for _ in ()).throw(OSError("down")), + ) + assert tmpl_mod._wait_tcp_down("127.0.0.1", 5986, 5.0) is True + + +def test_wait_tcp_down_still_up(monkeypatch: pytest.MonkeyPatch) -> None: + class _FakeConn: + def __enter__(self) -> _FakeConn: + return self + + def __exit__(self, *_: object) -> None: + pass + + monkeypatch.setattr(tmpl_mod.socket, "create_connection", lambda *a, **kw: _FakeConn()) + monkeypatch.setattr(tmpl_mod.time, "sleep", lambda _: None) + assert tmpl_mod._wait_tcp_down("127.0.0.1", 5986, 0.0) is False + + +def test_reboot_and_wait(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[str] = [] + transport = MagicMock() + transport.run.side_effect = lambda *_a, **_k: calls.append("reboot") + transport.close.side_effect = lambda: calls.append("close") + monkeypatch.setattr(tmpl_mod, "_wait_tcp_down", lambda *a, **k: True) + monkeypatch.setattr(tmpl_mod, "_wait_tcp", lambda *a, **k: True) + monkeypatch.setattr(tmpl_mod, "_wait_winrm", lambda *a, **k: calls.append("winrm")) + tmpl_mod._reboot_and_wait(transport, "10.0.0.5", 5986, 30) + assert calls == ["reboot", "close", "winrm"] + + +def test_reboot_and_wait_unreachable(monkeypatch: pytest.MonkeyPatch) -> None: + transport = MagicMock() + monkeypatch.setattr(tmpl_mod, "_wait_tcp_down", lambda *a, **k: True) + monkeypatch.setattr(tmpl_mod, "_wait_tcp", lambda *a, **k: False) + with pytest.raises(click.ClickException, match="not reachable after reboot"): + tmpl_mod._reboot_and_wait(transport, "10.0.0.5", 5986, 1) + + def test_wait_winrm_success() -> None: transport = MagicMock() transport.is_ready.return_value = True @@ -478,6 +519,7 @@ def _patch_prepare(monkeypatch: pytest.MonkeyPatch, backend: MagicMock, transpor monkeypatch.setattr(tmpl_mod, "WorkstationVmrunBackend", lambda **_: backend) monkeypatch.setattr(tmpl_mod, "WinRmTransport", lambda **_: transport) monkeypatch.setattr(tmpl_mod, "_wait_tcp", lambda *_: True) + monkeypatch.setattr(tmpl_mod, "_wait_tcp_down", lambda *_a, **_k: True) monkeypatch.setattr(tmpl_mod, "_wait_winrm", lambda *_: None) monkeypatch.setattr(tmpl_mod.time, "sleep", lambda _: None) monkeypatch.setattr(tmpl_mod.subprocess, "run", lambda *a, **_: None) @@ -640,6 +682,7 @@ def test_prepare_win_snapshot_exists_recreate( monkeypatch.setattr(tmpl_mod, "WorkstationVmrunBackend", lambda **_: backend) monkeypatch.setattr(tmpl_mod, "WinRmTransport", lambda **_: transport) monkeypatch.setattr(tmpl_mod, "_wait_tcp", lambda *_: True) + monkeypatch.setattr(tmpl_mod, "_wait_tcp_down", lambda *_a, **_k: True) monkeypatch.setattr(tmpl_mod, "_wait_winrm", lambda *_: None) monkeypatch.setattr(tmpl_mod.time, "sleep", lambda _: None) def _sp_run(cmd: list[str], **_: object) -> None: