tests/test_integration.py

"""Integration tests for the full GGUF splitting workflow"""

import os
import subprocess
import pathlib

import pytest

from src.gguf_utils import (
    get_llama_gguf_split_path,
    calculate_optimal_split_size,
    split_gguf_file,
)
from src.hf_utils import get_gguf_files_from_repo


class TestFullSplitWorkflow:
    """Integration tests for the complete splitting workflow"""

    def test_full_split_workflow_smollm2(self, temp_dir, test_gguf_url, test_gguf_filename):
        llama_path = get_llama_gguf_split_path()
        assert os.path.isfile(llama_path), f"llama-gguf-split not found at {llama_path}"
        
        gguf_path = temp_dir / test_gguf_filename
        print(f"\nDownloading {test_gguf_url}...")
        
        result = subprocess.run(
            ["curl", "-L", "-o", str(gguf_path), test_gguf_url],
            capture_output=True,
            text=True,
        )
        assert result.returncode == 0, f"Download failed: {result.stderr}"
        assert gguf_path.exists(), "Downloaded file does not exist"
        
        file_size_mb = gguf_path.stat().st_size / (1024 * 1024)
        print(f"Downloaded file size: {file_size_mb:.1f} MB")
        assert file_size_mb > 50, "Downloaded file is too small, might be corrupted"
        
        output_prefix = temp_dir / "gguf_split"
        max_size_mb = calculate_optimal_split_size(str(gguf_path), str(output_prefix))
        print(f"Calculated optimal split size: {max_size_mb} MB")
        assert max_size_mb > 0, "Split size should be positive"
        
        output_pattern = temp_dir / "gguf_split-"
        success = split_gguf_file(str(gguf_path), str(output_pattern), max_size_mb)
        assert success, "Splitting failed"
        
        split_files = sorted(temp_dir.glob("gguf_split-*.gguf"))
        print(f"Generated {len(split_files)} split files:")
        
        total_split_size = 0
        for f in split_files:
            size_mb = f.stat().st_size / (1024 * 1024)
            total_split_size += f.stat().st_size
            print(f"  - {f.name}: {size_mb:.1f} MB")
        
        assert len(split_files) >= 2, "Should have at least 2 split files"
        assert len(split_files) <= 10, "Should not have more than 10 split files for a 100MB model"
        
        for f in split_files:
            size_mb = f.stat().st_size / (1024 * 1024)
            assert size_mb <= max_size_mb * 1.1, f"Split file {f.name} exceeds max size"
        
        original_size = gguf_path.stat().st_size
        assert total_split_size >= original_size * 0.9, "Total split size is too small"
        assert total_split_size <= original_size * 1.5, "Total split size is too large"
        
        print(f"\nSplit verification passed!")
        print(f"  Original size: {original_size / (1024 * 1024):.1f} MB")
        print(f"  Total split size: {total_split_size / (1024 * 1024):.1f} MB")
        print(f"  Number of parts: {len(split_files)}")


class TestHuggingFaceIntegration:
    """Integration tests for Hugging Face API interactions"""

    def test_list_gguf_files_from_real_repo(self, test_repo_id):
        from huggingface_hub import HfApi
        
        api = HfApi()
        gguf_files = get_gguf_files_from_repo(test_repo_id, api)
        
        print(f"\nFound {len(gguf_files)} GGUF files in {test_repo_id}:")
        for f in gguf_files[:5]:
            print(f"  - {f}")
        
        assert len(gguf_files) > 0, f"No GGUF files found in {test_repo_id}"
        
        assert any("Q4_K_M" in f for f in gguf_files), "Test file Q4_K_M not found in repo"


class TestEnvironmentSetup:
    """Tests to verify the test environment is correctly set up"""

    def test_llama_gguf_split_available(self):
        path = get_llama_gguf_split_path()
        
        assert os.path.isfile(path), f"Binary not found at {path}"
        assert os.access(path, os.X_OK), f"Binary not executable at {path}"
        
        result = subprocess.run([path, "--help"], capture_output=True, text=True)
        assert result.returncode == 0, f"Binary failed to run: {result.stderr}"
        
        print(f"\nllama-gguf-split found at: {path}")
        print(f"Version info from --help (first 3 lines):")
        for line in result.stdout.split('\n')[:3]:
            print(f"  {line}")

    def test_curl_available(self):
        result = subprocess.run(["curl", "--version"], capture_output=True, text=True)
        assert result.returncode == 0, "curl not available"
        print(f"\ncurl version: {result.stdout.split(chr(10))[0]}")

    def test_python_dependencies(self):
        import gradio
        import huggingface_hub
        
        print(f"\ngradio version: {gradio.__version__}")
        print(f"huggingface_hub version: {huggingface_hub.__version__}")