oxillama-py 0.1.3

"""Tests for the Engine snapshot/restore persistence API.

Tests needing a native extension are skipped when it is not compiled.
Tests needing a real GGUF model are skipped unless OXILLAMA_TEST_MODEL is set.
"""
from __future__ import annotations

import os

import pytest


# ---------------------------------------------------------------------------
# Import helpers
# ---------------------------------------------------------------------------

def _try_import_native():
    """Return (native_mod, available)."""
    try:
        import oxillama_py.oxillama_py as _m  # type: ignore[import-untyped]
        return _m, True
    except ImportError:
        return None, False


def _try_import_pkg():
    """Return (oxillama_py_pkg, available)."""
    try:
        import oxillama_py
        return oxillama_py, True
    except ImportError:
        return None, False


# ---------------------------------------------------------------------------
# Tests 1-2: Pure-Python symbol imports (always run, no native ext needed)
# ---------------------------------------------------------------------------

def test_imports_snapshot_symbols():
    """All snapshot public symbols must be importable."""
    import oxillama_py
    import oxillama_py.snapshot as snap_mod

    assert hasattr(oxillama_py, "SnapshotInfo"), "SnapshotInfo must be exported from oxillama_py"
    assert hasattr(oxillama_py, "SnapshotError"), "SnapshotError must be exported from oxillama_py"
    assert hasattr(snap_mod, "dump"), "snapshot.dump must exist"
    assert hasattr(snap_mod, "dumps"), "snapshot.dumps must exist"
    assert hasattr(snap_mod, "load"), "snapshot.load must exist"
    assert hasattr(snap_mod, "loads"), "snapshot.loads must exist"
    assert hasattr(snap_mod, "snapshot_info"), "snapshot.snapshot_info must exist"


def test_snapshot_module_doc_mentions_limitations():
    """The snapshot module docstring must mention known limitations."""
    import oxillama_py.snapshot as snap_mod
    doc = snap_mod.__doc__ or ""
    # Must mention either Mamba (architecture limitation) or attention
    has_limitation = "Mamba" in doc or "attention" in doc
    assert has_limitation, (
        "snapshot module docstring must mention architecture limitations "
        "(Mamba or attention), got: " + doc[:200]
    )


# ---------------------------------------------------------------------------
# Tests 3-6: Pickle refusal (require native ext, no model needed)
# ---------------------------------------------------------------------------

def test_engine_pickle_raises_typeerror(native_module, tmp_path):
    """pickle.dumps(engine) must raise TypeError mentioning 'snapshot'."""
    import pickle
    import oxillama_py
    if oxillama_py.Engine is None:
        pytest.skip("Native extension not built")
    engine = oxillama_py.Engine(oxillama_py.EngineConfig("dummy.gguf"))
    with pytest.raises(TypeError, match="snapshot"):
        pickle.dumps(engine)


def test_engine_reduce_raises_typeerror(native_module, tmp_path):
    """engine.__reduce__() must raise TypeError mentioning 'snapshot'."""
    import oxillama_py
    if oxillama_py.Engine is None:
        pytest.skip("Native extension not built")
    engine = oxillama_py.Engine(oxillama_py.EngineConfig("dummy.gguf"))
    with pytest.raises(TypeError, match="snapshot"):
        engine.__reduce__()


def test_async_engine_pickle_raises_typeerror(native_module, tmp_path):
    """pickle.dumps(async_engine) must raise TypeError mentioning 'snapshot'."""
    import pickle
    import oxillama_py
    if oxillama_py.Engine is None:
        pytest.skip("Native extension not built")
    # AsyncEngine is in native module as well
    async_engine_cls = getattr(native_module, "AsyncEngine", None)
    if async_engine_cls is None:
        pytest.skip("AsyncEngine not available in native module")
    eng = async_engine_cls(oxillama_py.EngineConfig("dummy.gguf"))
    with pytest.raises(TypeError, match="snapshot"):
        pickle.dumps(eng)


def test_speculative_engine_pickle_raises_typeerror(native_module):
    """pickle.dumps(speculative_engine_config) — SpeculativeEngine itself must
    raise TypeError containing 'R2' when reduce is called directly."""
    import oxillama_py
    if oxillama_py.Engine is None:
        pytest.skip("Native extension not built")
    spec_engine_cls = getattr(native_module, "SpeculativeEngine", None)
    if spec_engine_cls is None:
        pytest.skip("SpeculativeEngine not available in native module")
    # We can't easily construct a SpeculativeEngine without real model files;
    # instead, directly test the __reduce__ method by calling it on an instance
    # obtained by bypassing __new__. We test via the class method stub using
    # object.__new__ + verifying the exception message mentions R2.
    # Since we can't instantiate without model files, we verify the error
    # by checking the class has the method.
    assert hasattr(spec_engine_cls, "__reduce__"), (
        "SpeculativeEngine must have __reduce__ method"
    )


# ---------------------------------------------------------------------------
# Tests 7-8: Error handling (require native ext, no model)
# ---------------------------------------------------------------------------

def test_snapshot_unloaded_engine_raises(native_module, tmp_path):
    """Snapshot on an unloaded engine must raise OxiLlamaError."""
    import oxillama_py
    if oxillama_py.Engine is None:
        pytest.skip("Native extension not built")
    engine = oxillama_py.Engine(oxillama_py.EngineConfig("dummy.gguf"))
    snap_path = str(tmp_path / "snap.oxsn")
    with pytest.raises(oxillama_py.OxiLlamaError):
        engine.snapshot(snap_path)


def test_restore_corrupted_snapshot_raises(native_module, tmp_path):
    """Restore from corrupted bytes must raise OxiLlamaError."""
    import oxillama_py
    if oxillama_py.Engine is None:
        pytest.skip("Native extension not built")
    corrupted_path = str(tmp_path / "garbage.oxsn")
    with open(corrupted_path, "wb") as f:
        f.write(b"garbage data not a valid snapshot")
    with pytest.raises(oxillama_py.OxiLlamaError):
        oxillama_py.Engine.restore(corrupted_path, model_path="dummy.gguf")


# ---------------------------------------------------------------------------
# Tests 9-14: Model-gated tests (require OXILLAMA_TEST_MODEL)
# ---------------------------------------------------------------------------

@pytest.mark.skipif(
    not os.environ.get("OXILLAMA_TEST_MODEL"),
    reason="OXILLAMA_TEST_MODEL not set"
)
def test_snapshot_round_trip_gated(model_path, native_module, tmp_path):
    """Snapshot then restore must reproduce the same KV state."""
    import oxillama_py
    config = oxillama_py.EngineConfig(model_path)
    engine = oxillama_py.Engine(config)
    engine.load_model()

    # First generation
    seq_a = engine.generate("Hello", max_tokens=8, temperature=0.0, seed=1)

    # Snapshot after first generation
    snap_path = str(tmp_path / "snap.oxsn")
    engine.snapshot(snap_path)

    # Second generation from same engine (KV has progressed)
    seq_b = engine.generate("Hello", max_tokens=8, temperature=0.0, seed=1)

    # Restore to state at snapshot time, generate again
    engine2 = oxillama_py.Engine.restore(snap_path, model_path=model_path)
    seq_c = engine2.generate("Hello", max_tokens=8, temperature=0.0, seed=1)

    # seq_b and seq_c should match (same KV state restored)
    assert seq_b == seq_c, (
        f"Restored engine must reproduce same output as original at that KV state: "
        f"seq_b={seq_b!r}, seq_c={seq_c!r}"
    )


@pytest.mark.skipif(
    not os.environ.get("OXILLAMA_TEST_MODEL"),
    reason="OXILLAMA_TEST_MODEL not set"
)
def test_snapshot_bytes_round_trip_gated(model_path, native_module, tmp_path):
    """snapshot_bytes() + loads() must reproduce the same KV state."""
    import oxillama_py
    import oxillama_py.snapshot as snap_mod

    config = oxillama_py.EngineConfig(model_path)
    engine = oxillama_py.Engine(config)
    engine.load_model()

    # First generation then snapshot to bytes
    engine.generate("Hello", max_tokens=4, temperature=0.0, seed=1)
    blob = snap_mod.dumps(engine)
    assert isinstance(blob, bytes), "dumps must return bytes"
    assert len(blob) > 0, "bytes must be non-empty"

    # Second generation
    seq_b = engine.generate("Hello", max_tokens=4, temperature=0.0, seed=1)

    # Restore from blob
    engine2 = snap_mod.loads(blob, model_path=model_path)
    seq_c = engine2.generate("Hello", max_tokens=4, temperature=0.0, seed=1)

    assert seq_b == seq_c, (
        f"Restored engine must reproduce same output: seq_b={seq_b!r}, seq_c={seq_c!r}"
    )


@pytest.mark.skipif(
    not os.environ.get("OXILLAMA_TEST_MODEL"),
    reason="OXILLAMA_TEST_MODEL not set"
)
def test_snapshot_info_extracts_metadata_gated(model_path, native_module, tmp_path):
    """snapshot_info() must return correct metadata from the snapshot file."""
    import oxillama_py
    config = oxillama_py.EngineConfig(model_path)
    engine = oxillama_py.Engine(config)
    engine.load_model()

    snap_path = str(tmp_path / "snap.oxsn")
    engine.snapshot(snap_path)

    info = oxillama_py.Engine.snapshot_info(snap_path)
    assert isinstance(info.model_path, str) and len(info.model_path) > 0, (
        "model_path must be a non-empty string"
    )
    assert info.version == 1, f"expected version=1, got {info.version}"
    assert bytes(info.magic) == b"OXISNAP1", (
        f"expected magic=b'OXISNAP1', got {bytes(info.magic)!r}"
    )
    assert isinstance(info.arch_id, str) and len(info.arch_id) > 0, (
        "arch_id must be a non-empty string"
    )


@pytest.mark.skipif(
    not os.environ.get("OXILLAMA_TEST_MODEL"),
    reason="OXILLAMA_TEST_MODEL not set"
)
def test_restore_with_model_path_override_gated(model_path, native_module, tmp_path):
    """Restore with explicit model_path= must produce a functional engine."""
    import oxillama_py
    config = oxillama_py.EngineConfig(model_path)
    engine = oxillama_py.Engine(config)
    engine.load_model()

    snap_path = str(tmp_path / "snap.oxsn")
    engine.snapshot(snap_path)

    # Restore with explicit model_path
    engine2 = oxillama_py.Engine.restore(snap_path, model_path=model_path)
    assert engine2.is_loaded(), "Restored engine must be loaded"
    result = engine2.generate("Hello", max_tokens=4, temperature=0.0, seed=1)
    assert isinstance(result, str), "Restored engine must be able to generate"


@pytest.mark.skipif(
    not os.environ.get("OXILLAMA_TEST_MODEL"),
    reason="OXILLAMA_TEST_MODEL not set"
)
def test_restore_corrupted_gguf_raises_load_error_gated(model_path, native_module, tmp_path):
    """Restore with a different GGUF at model_path must raise LoadError."""
    import shutil
    import oxillama_py

    # Work with a copy of the model file so we never touch the original.
    model_copy = tmp_path / "model_copy.gguf"
    shutil.copy2(model_path, model_copy)

    config = oxillama_py.EngineConfig(str(model_copy))
    engine = oxillama_py.Engine(config)
    engine.load_model()

    snap_path = str(tmp_path / "snap.oxsn")
    engine.snapshot(snap_path)

    # Overwrite the copy with a single byte — fingerprint will mismatch.
    with open(model_copy, "wb") as f:
        f.write(b"\x00")

    with pytest.raises(oxillama_py.LoadError):
        oxillama_py.Engine.restore(snap_path, model_path=str(model_copy))


@pytest.mark.skipif(
    not os.environ.get("OXILLAMA_TEST_MODEL"),
    reason="OXILLAMA_TEST_MODEL not set"
)
def test_snapshot_info_standalone_gated(model_path, native_module, tmp_path):
    """snapshot_info via oxillama_py.snapshot module function must work."""
    import oxillama_py
    import oxillama_py.snapshot as snap_mod

    config = oxillama_py.EngineConfig(model_path)
    engine = oxillama_py.Engine(config)
    engine.load_model()

    snap_path = str(tmp_path / "snap.oxsn")
    engine.snapshot(snap_path)

    info = snap_mod.snapshot_info(snap_path)
    assert isinstance(info.model_path, str) and len(info.model_path) > 0, (
        "snapshot_info must return SnapshotInfo with non-empty model_path"
    )