unsloth/tests/python/test_windows_on_arm_python_arch.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.

"""Windows on ARM must end on an x64 interpreter or stop, since PyPI has no win_arm64 pyarrow wheel."""

from __future__ import annotations

import os
import re
import shutil
from pathlib import Path

import pytest

from unsloth_pwsh_runner import run_pwsh


REPO_ROOT = Path(__file__).resolve().parents[2]
INSTALL_PS1 = REPO_ROOT / "install.ps1"
POWERSHELLS = [shell for shell in ("pwsh", "powershell") if shutil.which(shell)]


def _function(name: str) -> str:
    """Extracts a function from install.ps1 by indented braces, so rewording a comment cannot change it."""
    source = INSTALL_PS1.read_text(encoding = "utf-8")
    match = re.search(rf"    function {name} \{{.*?\n    \}}\n", source, flags = re.DOTALL)
    assert match is not None, f"install.ps1 no longer defines {name}"
    return match.group(0)


def _run(
    shell: str,
    script: str,
    env: dict[str, str] | None = None,
) -> str:
    # run_pwsh: both shells write UTF-8 while `text = True` decodes with the locale codec.
    result = run_pwsh(
        [shell, "-NoProfile", "-NonInteractive", "-Command", script],
        check = True,
        capture_output = True,
        text = True,
        encoding = "utf-8",
        errors = "strict",
        env = {**os.environ, **(env or {})},
        timeout = 60,
    )
    return result.stdout.strip()


# The shared preamble: the two host helpers the functions under test call, stubbed so the
# decision is driven by the case rather than by the machine running the suite. Neither stub
# stands in for logic under test here: Get-HostMachineArch is exercised against real
# hardware on the windows-11-arm leg, and Get-PythonPlatformTag asks an interpreter for
# sysconfig.get_platform(), which needs a Windows interpreter to answer win-arm64.
def _preamble(host_arch: str, venv_tag: str) -> str:
    # Verbatim from install.ps1, not a stub: both the fresh-selection path and the venv
    # re-check call it, so a change to what the opt-out accepts has to move these cases too.
    opt_out = _function("Test-Arm64PythonOptOut")
    return f"""
$ErrorActionPreference = "Stop"
$script:StudioVtOk = $false
$script:StudioStdoutRedirected = $false
function Write-StudioLine {{
    param([string]$Message = "", [string]$ForegroundColor)
    Write-Host $Message
}}
function step {{ param([string]$Label, [string]$Value, [string]$Color = "Green") Write-Host "STEP ${{Label}}: $Value" }}
function substep {{ param([string]$Message, [string]$Color = "DarkGray") Write-Host "SUBSTEP $Message" }}
function Get-HostMachineArch {{ return "{host_arch}" }}
function Get-PythonPlatformTag {{ param([string]$Exe) return "{venv_tag}" }}
{opt_out}
"""


@pytest.mark.skipif(not POWERSHELLS, reason = "PowerShell is unavailable")
@pytest.mark.parametrize("shell", POWERSHELLS)
@pytest.mark.parametrize(
    "host_arch,venv_tag,selected_arch,venv_exists,expected",
    [
        ("arm64", "win-arm64", "x86_64", True, True),
        ("arm64", "win-amd64", "x86_64", True, False),
        ("x86_64", "win-amd64", "x86_64", True, False),
        # Must NOT fire: native-ARM64-wheel hosts look like this and a rebuild would regress them.
        ("arm64", "win-arm64", "arm64", True, False),
        ("arm64", "win-arm64", "x86_64", False, False),
        ("arm64", "", "x86_64", True, False),
    ],
)
def test_a_reused_environment_is_re_checked_for_its_interpreter_architecture(
    tmp_path: Path,
    shell: str,
    host_arch: str,
    venv_tag: str,
    selected_arch: str,
    venv_exists: bool,
    expected: bool,
):
    """The swap only covers a fresh selection, so reuse has to be re-checked separately."""
    venv_python = tmp_path / "python.exe"
    if venv_exists:
        venv_python.write_text("", encoding = "utf-8")
    selected = f'@{{ Version = "3.13"; Path = "C:\\p\\python.exe"; Arch = "{selected_arch}" }}'
    script = (
        _preamble(host_arch, venv_tag)
        + f"""
{_function("Test-StudioVenvArchMismatch")}
$selected = {selected}
$mismatch = Test-StudioVenvArchMismatch -VenvPython "{venv_python}" -SelectedPython $selected
Write-Host ("MISMATCH=" + $mismatch)
"""
    )
    out = _run(shell, script)
    assert f"MISMATCH={expected}" in out


@pytest.mark.skipif(not POWERSHELLS, reason = "no PowerShell on this host")
@pytest.mark.parametrize("shell", POWERSHELLS)
@pytest.mark.parametrize("value", ["1", "true", "YES", "on"])
def test_the_opt_out_also_protects_an_existing_arm64_environment(
    tmp_path: Path, shell: str, value: str
):
    """The opt-out has to be read HERE, not only where a fresh interpreter is chosen.

    Resolve-WindowsOnArmX64Python never runs when the ordinary probe already found x64, so
    on an ARM64 box with both interpreters the variable was never consulted and this branch
    moved aside the very environment the user asked to keep. The rebuild is not a migration:
    nothing installed into the old venv carries over.
    """
    venv_python = tmp_path / "python.exe"
    venv_python.write_text("", encoding = "utf-8")
    selected = '@{ Version = "3.13"; Path = "C:\\p\\python.exe"; Arch = "x86_64" }'
    script = (
        _preamble("arm64", "win-arm64")
        + f"""
{_function("Test-StudioVenvArchMismatch")}
$selected = {selected}
$mismatch = Test-StudioVenvArchMismatch -VenvPython "{venv_python}" -SelectedPython $selected
Write-Host ("MISMATCH=" + $mismatch)
"""
    )
    out = _run(shell, script, env = {"UNSLOTH_ALLOW_ARM64_PYTHON": value})
    assert (
        "MISMATCH=False" in out
    ), "the opt-out was ignored and an existing ARM64 environment was scheduled for replacement"


@pytest.mark.skipif(not POWERSHELLS, reason = "no PowerShell on this host")
@pytest.mark.parametrize("shell", POWERSHELLS)
def test_the_arm64_environment_is_still_replaced_without_the_opt_out(tmp_path: Path, shell: str):
    """The other direction, so the guard above cannot be satisfied by never firing."""
    venv_python = tmp_path / "python.exe"
    venv_python.write_text("", encoding = "utf-8")
    selected = '@{ Version = "3.13"; Path = "C:\\p\\python.exe"; Arch = "x86_64" }'
    script = (
        _preamble("arm64", "win-arm64")
        + f"""
{_function("Test-StudioVenvArchMismatch")}
$selected = {selected}
$mismatch = Test-StudioVenvArchMismatch -VenvPython "{venv_python}" -SelectedPython $selected
Write-Host ("MISMATCH=" + $mismatch)
"""
    )
    out = _run(shell, script)
    assert "MISMATCH=True" in out


def test_both_arm64_opt_out_readers_go_through_one_helper():
    """Two spellings of the same question drift. The fresh-selection path and the reuse
    re-check must read it the same way, or an accepted value works in one and not the
    other."""
    source = INSTALL_PS1.read_text(encoding = "utf-8")
    assert source.count("function Test-Arm64PythonOptOut") == 1
    assert (
        source.count("$env:UNSLOTH_ALLOW_ARM64_PYTHON -in") == 1
    ), "the opt-out is compared in more than one place"
    assert (
        source.count("Test-Arm64PythonOptOut") >= 3
    ), "both the selection path and the venv re-check have to call it"


def test_the_arm64_rebuild_tells_the_user_what_it_does_not_carry_over():
    """Arch rebuilds unavoidably drop user packages; the branch must name the old tree and how to opt out"""
    source = INSTALL_PS1.read_text(encoding = "utf-8")
    start = source.find("Test-StudioVenvArchMismatch -VenvPython")
    assert start != -1
    block = source[start : source.find('step "venv" "creating Python', start)]
    assert "unsloth_studio.arm64" in block, "the rebuild does not say where the old venv went"
    assert "UNSLOTH_ALLOW_ARM64_PYTHON" in block, "the rebuild does not name its opt-out"


def test_the_venv_reuse_path_re_checks_before_it_reuses():
    """The arch re-check must run before the migrated-environment reuse branch, not after it commits."""
    source = INSTALL_PS1.read_text(encoding = "utf-8")
    recheck = source.find("Test-StudioVenvArchMismatch -VenvPython")
    create = source.find('step "venv" "creating Python')
    reuse = source.find('step "venv" "using migrated environment"')
    assert recheck != -1, "install.ps1 no longer re-checks the reused interpreter"
    assert create != -1 and reuse != -1
    assert recheck < create < reuse, "the re-check must precede the create/reuse branch"
    # $_Migrated must be cleared, or --no-deps installs nothing into the fresh venv.
    tail = source[recheck:create]
    assert "$_Migrated = $false" in tail


# ── The x64 bootstrap has to respect conda too ──
# ── The environment the rebuild promises to keep has to still be there ──


def _rollback_preamble(studio_home: Path) -> str:
    return f"""
$ErrorActionPreference = "Stop"
function Write-StudioLine {{
    param([string]$Message = "", [string]$ForegroundColor)
    Write-Host $Message
}}
function substep {{ param([string]$Message, [string]$Color = "DarkGray") Write-Host "SUBSTEP $Message" }}
$StudioHome = "{studio_home}"
function Remove-StudioVenvTreeWithRetry {{
    param([string]$Path, [string]$Label)
    Write-Host "REMOVED $Path"
    Remove-Item -LiteralPath $Path -Recurse -Force
    return $true
}}
{_function("Test-StudioPathPresent")}
{_function("Complete-StudioVenvRollback")}
"""


@pytest.mark.skipif(not POWERSHELLS, reason = "PowerShell is unavailable")
@pytest.mark.parametrize("shell", POWERSHELLS)
def test_the_preserved_arm64_environment_survives_a_successful_install(tmp_path: Path, shell: str):
    """The rebuild promises the ARM64 environment is kept so packages can be copied out.

    Complete-StudioVenvRollback deletes the backup on success, which is right for every
    other rollback but wrong for this one: the old tree holds packages the new architecture
    cannot have, and deleting it discards them and the message that said otherwise.
    """
    home = tmp_path / "studio"
    backup = home / "unsloth_studio.rollback.20260101000000.4242"
    (backup / "Lib").mkdir(parents = True)
    (backup / "Lib" / "marker.txt").write_text("a package the user added", encoding = "utf-8")

    script = (
        _rollback_preamble(home)
        + f"""
$script:StudioVenvRollbackActive = $true
$script:StudioVenvRollbackDir = "{backup}"
$script:StudioVenvRollbackPartial = $false
$script:StudioVenvRollbackPreserve = $true
Complete-StudioVenvRollback
Write-Host ("COMMITTED=" + $script:StudioInstallCommitted)
Write-Host ("PRESERVE=" + $script:StudioVenvRollbackPreserve)
"""
    )
    out = _run(shell, script)
    assert "COMMITTED=True" in out, "the commit flag is what stops a later restore"
    assert "PRESERVE=False" in out, "the flag must not survive into a later rollback"
    assert "REMOVED" not in out, "the tree the user was told to read must not be deleted"
    assert not backup.exists(), "it has to leave the rollback namespace, see below"

    kept = [path for path in home.iterdir() if path.name.startswith("unsloth_studio.arm64.")]
    assert len(kept) == 1, f"expected one preserved ARM64 tree, found {[p.name for p in kept]}"
    assert (kept[0] / "Lib" / "marker.txt").read_text(encoding = "utf-8") == (
        "a package the user added"
    )
    assert (
        f"previous ARM64 environment kept at {kept[0]}" in out
    ), "a preserved tree nobody is told the path of is just disk usage"


@pytest.mark.skipif(not POWERSHELLS, reason = "PowerShell is unavailable")
@pytest.mark.parametrize("shell", POWERSHELLS)
def test_an_ordinary_rollback_is_still_deleted_on_success(tmp_path: Path, shell: str):
    """The preserve flag is opt-in, and every other rollback keeps the old behaviour.

    Without this the fix trades one bug for a slow leak: a full environment left behind by
    every reinstall, on the machine of every user who never asked for one.
    """
    home = tmp_path / "studio"
    backup = home / "unsloth_studio.rollback.20260101000000.4242"
    backup.mkdir(parents = True)

    script = (
        _rollback_preamble(home)
        + f"""
$script:StudioVenvRollbackActive = $true
$script:StudioVenvRollbackDir = "{backup}"
$script:StudioVenvRollbackPartial = $false
$script:StudioVenvRollbackPreserve = $false
Complete-StudioVenvRollback
"""
    )
    out = _run(shell, script)
    assert f"REMOVED {backup}" in out
    assert not backup.exists()
    assert not list(home.iterdir()), "nothing is kept when nothing was promised"


def test_the_preserved_tree_leaves_the_swept_rollback_namespace():
    """Remove-StaleStudioVenvRollbacks globs unsloth_studio.rollback.*, and preserves an entry
    only while its owning PID is alive."""
    source = INSTALL_PS1.read_text(encoding = "utf-8")
    complete = source.index("function Complete-StudioVenvRollback")
    block = source[complete : source.index("\n    }\n", complete)]
    # Comments stripped: they mention the rollback name while explaining why it is not used.
    code = "\n".join(line for line in block.splitlines() if not line.lstrip().startswith("#"))
    assert "unsloth_studio.arm64." in code, "the preserved copy is not renamed"
    assert (
        "unsloth_studio.rollback." not in code
    ), "a preserved copy under the swept name is removed by the next install"
    sweep = source.index("function Remove-StaleStudioVenvRollbacks")
    sweep_block = source[sweep : source.index("\n    }\n", sweep)]
    assert (
        "'unsloth_studio.rollback.*'" in sweep_block
    ), "the sweep no longer matches what this test reasons about"
    assert "unsloth_studio.arm64" not in sweep_block


# The reinstall moves the venv to its rollback path before x64 selection, so the platform
# tag must be read before that move.


def test_the_platform_tag_is_read_before_the_rollback_move():
    source = INSTALL_PS1.read_text(encoding = "utf-8")
    record = source.find("$script:PrevVenvPlatformTag = Get-PythonPlatformTag")
    move = source.find("Start-StudioVenvRollback -ExistingDir $VenvDir")
    recheck = source.find("Test-StudioVenvArchMismatch -VenvPython")
    assert record != -1, "install.ps1 no longer records the existing interpreter's tag"
    assert record < move, "the tag is read after the move that takes the interpreter away"
    assert (
        "-RecordedTag $script:PrevVenvPlatformTag" in source[recheck : recheck + 400]
    ), "the re-check does not use the recorded tag, so it cannot see a moved environment"


def _arch_mismatch_preamble(
    host_arch: str, opt_out: bool, venv_python_exists: bool, live_tag: str, recorded_tag: str
) -> str:
    selected = '@{ Version = "3.13"; Arch = "x86_64" }'
    exists = "$true" if venv_python_exists else "$false"
    recorded = f'"{recorded_tag}"' if recorded_tag else "$null"
    return f"""
$ErrorActionPreference = "Stop"
function substep {{ param([string]$Message = "", [string]$Color) Write-Host "SUBSTEP $Message" }}
function Get-HostMachineArch {{ return "{host_arch}" }}
function Test-Arm64PythonOptOut {{ return ${str(opt_out).lower()} }}
function Get-PythonPlatformTag {{ param($Exe) return "{live_tag}" }}
function Test-Path {{ param([Parameter(ValueFromRemainingArguments = $true)]$Rest) return {exists} }}
{_function("Test-StudioVenvArchMismatch")}
$answer = Test-StudioVenvArchMismatch -VenvPython "C:\\venv\\Scripts\\python.exe" `
    -SelectedPython {selected} -RecordedTag {recorded}
Write-Host ("ANSWER=" + $answer)
"""


@pytest.mark.skipif(not POWERSHELLS, reason = "PowerShell is unavailable")
@pytest.mark.parametrize("shell", POWERSHELLS)
def test_a_moved_arm64_environment_is_still_detected(shell: str):
    """The case the reinstall actually takes: the interpreter has been moved aside already."""
    out = _run(
        shell,
        _arch_mismatch_preamble("arm64", False, False, "win-amd64", "win-arm64"),
    )
    assert out.splitlines()[-1] == "ANSWER=True", out


@pytest.mark.skipif(not POWERSHELLS, reason = "PowerShell is unavailable")
@pytest.mark.parametrize("shell", POWERSHELLS)
def test_a_moved_x64_environment_is_not_mistaken_for_one(shell: str):
    """The other half. A recorded tag that is not win-arm64 answers no, and so does no
    recorded tag at all, which is the fresh-install case."""
    out = _run(shell, _arch_mismatch_preamble("arm64", False, False, "win-arm64", "win-amd64"))
    assert out.splitlines()[-1] == "ANSWER=False", out
    out = _run(shell, _arch_mismatch_preamble("arm64", False, False, "win-arm64", ""))
    assert out.splitlines()[-1] == "ANSWER=False", out


@pytest.mark.skipif(not POWERSHELLS, reason = "PowerShell is unavailable")
@pytest.mark.parametrize("shell", POWERSHELLS)
def test_the_opt_out_message_is_only_printed_for_an_arm64_environment(shell: str):
    """`UNSLOTH_ALLOW_ARM64_PYTHON=1` on an x64 environment used to announce that a native
    ARM64 environment was being kept, which was true of nothing on the machine."""
    out = _run(shell, _arch_mismatch_preamble("arm64", True, False, "win-amd64", "win-amd64"))
    assert "SUBSTEP" not in out, out
    assert out.splitlines()[-1] == "ANSWER=False", out

    out = _run(shell, _arch_mismatch_preamble("arm64", True, True, "win-arm64", "win-arm64"))
    assert "UNSLOTH_ALLOW_ARM64_PYTHON" in out, out
    assert out.splitlines()[-1] == "ANSWER=False", out


def test_the_architecture_rebuild_does_not_move_the_venv_twice():
    """A reinstall already moved $VenvDir to rollback; a second Start-StudioVenvRollback would throw."""
    source = INSTALL_PS1.read_text(encoding = "utf-8")
    branch = source.index("windows on arm: the existing environment runs native ARM64 Python")
    # Search forward from the branch rather than a fixed window that breaks as the chain grows.
    guard = source.index("if ($script:StudioVenvRollbackActive)", branch)
    move = source.index("Start-StudioVenvRollback -ExistingDir $VenvDir", branch)
    assert guard < move, "the rebuild still moves the environment unconditionally"
    assert (
        "$script:StudioVenvRollbackPreserve = $true" in source[guard:move]
    ), "an already-active rollback is not marked for preservation"


def test_the_opt_out_keeps_the_environment_it_says_it_keeps():
    """`UNSLOTH_ALLOW_ARM64_PYTHON=1` returning $false is not, by itself, keeping anything:
    on a host that already has an x64 Python the reinstall has moved the ARM64 environment
    into a rollback that success deletes. The opt-out marks that rollback preserved."""
    source = INSTALL_PS1.read_text(encoding = "utf-8")
    optout = source.index(
        '$script:PrevVenvPlatformTag -eq "win-arm64" -and (Test-Arm64PythonOptOut)'
    )
    create = source.index('step "venv" "creating Python')
    assert optout < create, "the opt-out preservation runs after the rebuild has started"
    assert "$script:StudioVenvRollbackPreserve = $true" in source[optout : optout + 600]
    assert (
        "unsloth_studio.arm64.*" in source[optout : optout + 600]
    ), "the user is not told where the kept environment is"