neuron-tai/tests/test_split_gguf_provision.py

"""DGR-026 — resumable, hash-verifying split-GGUF provisioning to mounted-drive storage.

Deterministic, offline, GPU-free, and download-free: every split here is a
tiny local fixture file; nothing is downloaded from a network.
"""

from __future__ import annotations

import hashlib

import pytest

from meshnet_node.split_gguf.manifest import parse_split_artifact_manifest
from meshnet_node.split_gguf.provision import (
    SplitProvisionError,
    local_directory_fetcher,
    provision_split_artifact,
    reject_home_path,
    verify_provisioned_split_artifact,
)


def _sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def _rev(label: str) -> str:
    return hashlib.sha1(label.encode()).hexdigest()


SPLIT_A = b"deepseek-v4-flash split A payload bytes " * 100
SPLIT_B = b"deepseek-v4-flash split B payload bytes, a bit longer than A " * 130


def _manifest_doc() -> dict:
    return {
        "schema_version": 1,
        "manifest_id": "deepseek-v4-flash-fixture",
        "manifest_version": "test.1",
        "quantization": "Q4_K_M",
        "source": {
            "artifact_id": "deepseek-v4-flash",
            "repo_id": "example/deepseek-v4-flash-gguf",
            "revision": _rev("source"),
            "sha256": _sha256_bytes(b"whole-model"),
            "size_bytes": len(SPLIT_A) + len(SPLIT_B),
        },
        "tokenizer": {
            "repo_id": "example/deepseek-v4-flash",
            "revision": _rev("tokenizer"),
            "sha256": _sha256_bytes(b"tokenizer"),
        },
        "total_bytes": len(SPLIT_A) + len(SPLIT_B),
        "splits": [
            {
                "name": "split-a.gguf",
                "size_bytes": len(SPLIT_A),
                "sha256": _sha256_bytes(SPLIT_A),
                "role": "layers-0-20",
                "shard_start": 0,
                "shard_end": 20,
            },
            {
                "name": "split-b.gguf",
                "size_bytes": len(SPLIT_B),
                "sha256": _sha256_bytes(SPLIT_B),
                "role": "layers-20-43",
                "shard_start": 20,
                "shard_end": 43,
            },
        ],
    }


@pytest.fixture
def manifest():
    return parse_split_artifact_manifest(_manifest_doc(), origin="<fixture>")


@pytest.fixture
def source_dir(tmp_path):
    d = tmp_path / "source"
    d.mkdir()
    (d / "split-a.gguf").write_bytes(SPLIT_A)
    (d / "split-b.gguf").write_bytes(SPLIT_B)
    return d


# --------------------------------------------------------------------------
# /home rejection
# --------------------------------------------------------------------------


def test_provisioning_refuses_a_destination_under_home(manifest, source_dir):
    with pytest.raises(SplitProvisionError, match="never /home"):
        provision_split_artifact(manifest, "/home/someuser/models", local_directory_fetcher(source_dir))


def test_reject_home_path_refuses_home_itself():
    with pytest.raises(SplitProvisionError, match="never /home"):
        reject_home_path("/home")


def test_reject_home_path_refuses_a_nested_home_subdirectory():
    with pytest.raises(SplitProvisionError, match="never /home"):
        reject_home_path("/home/someuser/.cache/meshnet/models")


def test_reject_home_path_accepts_a_mounted_drive_path(tmp_path):
    dest = tmp_path / "mnt" / "models"
    assert reject_home_path(dest) == dest.expanduser().resolve()


def test_verify_provisioned_also_refuses_home(manifest):
    with pytest.raises(SplitProvisionError, match="never /home"):
        verify_provisioned_split_artifact(manifest, "/home/someuser/models")


# --------------------------------------------------------------------------
# Happy path + idempotent re-run
# --------------------------------------------------------------------------


def test_provisioning_fetches_and_verifies_every_split(manifest, source_dir, tmp_path):
    dest = tmp_path / "mnt" / "models"
    result = provision_split_artifact(manifest, dest, local_directory_fetcher(source_dir))

    assert result.dest_dir == dest.resolve()
    assert set(result.verified_splits) == {"split-a.gguf", "split-b.gguf"}
    assert (dest / "split-a.gguf").read_bytes() == SPLIT_A
    assert (dest / "split-b.gguf").read_bytes() == SPLIT_B
    assert not (dest / "split-a.gguf.partial").exists()
    assert not (dest / "split-b.gguf.partial").exists()

    verify_provisioned_split_artifact(manifest, dest)  # does not raise


def test_a_second_provisioning_run_is_a_no_op_over_complete_splits(manifest, source_dir, tmp_path):
    dest = tmp_path / "mnt" / "models"
    provision_split_artifact(manifest, dest, local_directory_fetcher(source_dir))

    # Delete the source so a second run could not possibly re-fetch anything;
    # the already-complete, hash-correct splits must be recognized as done.
    (source_dir / "split-a.gguf").unlink()
    (source_dir / "split-b.gguf").unlink()

    result = provision_split_artifact(manifest, dest, local_directory_fetcher(source_dir))
    assert set(result.verified_splits) == {"split-a.gguf", "split-b.gguf"}


# --------------------------------------------------------------------------
# Interrupted download → resume
# --------------------------------------------------------------------------


def test_an_interrupted_partial_download_resumes_from_its_exact_byte_offset(manifest, source_dir, tmp_path):
    dest = tmp_path / "mnt" / "models"
    dest.mkdir(parents=True)

    # Simulate an interrupted prior attempt: split-a is half-written as a
    # `.partial` file with correct bytes so-far; split-b has not started.
    cut = len(SPLIT_A) // 2
    (dest / "split-a.gguf.partial").write_bytes(SPLIT_A[:cut])

    calls: list[tuple[str, int]] = []
    real_fetch = local_directory_fetcher(source_dir)

    def tracking_fetch(split, dest_path, resume_from_bytes):
        calls.append((split.name, resume_from_bytes))
        real_fetch(split, dest_path, resume_from_bytes)

    result = provision_split_artifact(manifest, dest, tracking_fetch)

    assert ("split-a.gguf", cut) in calls  # resumed from the exact offset, not from 0
    assert ("split-b.gguf", 0) in calls
    assert (dest / "split-a.gguf").read_bytes() == SPLIT_A
    assert (dest / "split-b.gguf").read_bytes() == SPLIT_B
    assert set(result.verified_splits) == {"split-a.gguf", "split-b.gguf"}


def test_a_partial_larger_than_the_manifest_size_is_discarded_and_restarted(manifest, source_dir, tmp_path):
    dest = tmp_path / "mnt" / "models"
    dest.mkdir(parents=True)
    (dest / "split-a.gguf.partial").write_bytes(SPLIT_A + b"stray corrupt trailing bytes")

    result = provision_split_artifact(manifest, dest, local_directory_fetcher(source_dir))

    assert (dest / "split-a.gguf").read_bytes() == SPLIT_A
    assert set(result.verified_splits) == {"split-a.gguf", "split-b.gguf"}


# --------------------------------------------------------------------------
# Missing split
# --------------------------------------------------------------------------


def test_a_missing_split_source_file_raises(manifest, source_dir, tmp_path):
    (source_dir / "split-b.gguf").unlink()
    dest = tmp_path / "mnt" / "models"

    with pytest.raises(SplitProvisionError, match="split source is missing"):
        provision_split_artifact(manifest, dest, local_directory_fetcher(source_dir))


def test_verify_reports_a_split_missing_from_the_destination(manifest, source_dir, tmp_path):
    dest = tmp_path / "mnt" / "models"
    provision_split_artifact(manifest, dest, local_directory_fetcher(source_dir))
    (dest / "split-b.gguf").unlink()

    with pytest.raises(SplitProvisionError, match="missing split"):
        verify_provisioned_split_artifact(manifest, dest)


# --------------------------------------------------------------------------
# Hash mismatch
# --------------------------------------------------------------------------


def test_a_hash_mismatched_source_file_is_rejected_and_not_left_on_disk(manifest, source_dir, tmp_path):
    # Same size as the pinned split so the mismatch is caught by hash, not by
    # the incomplete-byte-count check.
    (source_dir / "split-a.gguf").write_bytes(b"x" * len(SPLIT_A))
    dest = tmp_path / "mnt" / "models"

    with pytest.raises(SplitProvisionError, match="hash mismatch"):
        provision_split_artifact(manifest, dest, local_directory_fetcher(source_dir))

    assert not (dest / "split-a.gguf").exists()
    assert not (dest / "split-a.gguf.partial").exists()


def test_a_destination_file_with_wrong_hash_is_not_trusted_and_is_replaced(manifest, source_dir, tmp_path):
    dest = tmp_path / "mnt" / "models"
    dest.mkdir(parents=True)
    (dest / "split-a.gguf").write_bytes(b"x" * len(SPLIT_A))  # right size, wrong content

    result = provision_split_artifact(manifest, dest, local_directory_fetcher(source_dir))

    assert (dest / "split-a.gguf").read_bytes() == SPLIT_A
    assert set(result.verified_splits) == {"split-a.gguf", "split-b.gguf"}


def test_verify_reports_a_hash_mismatched_destination_file(manifest, source_dir, tmp_path):
    dest = tmp_path / "mnt" / "models"
    provision_split_artifact(manifest, dest, local_directory_fetcher(source_dir))
    (dest / "split-a.gguf").write_bytes(b"corrupted after the fact" + SPLIT_A)

    with pytest.raises(SplitProvisionError, match="mismatch"):
        verify_provisioned_split_artifact(manifest, dest)