tests/test_processing.py

"""Unit tests for processing module"""

import pathlib
import subprocess
import sys
import tempfile
from unittest.mock import Mock, patch

import gradio as gr
import pytest


sys.path.insert(0, "src")

from src.processing import (
    save_locally,
    upload_to_hf,
    _download_with_retry,
    _validate_download,
    _generate_readme_content,
    _cleanup_temp_files,
)


class TestValidateDownload:
    """Test _validate_download function"""

    def test_valid_file(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = pathlib.Path(tmp_dir)
            gguf_path = tmp_path / "model.gguf"
            gguf_path.write_text("dummy content")

            _validate_download(gguf_path)

    def test_missing_file(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = pathlib.Path(tmp_dir)
            gguf_path = tmp_path / "nonexistent.gguf"

            with pytest.raises(Exception, match="Download failed - file not found"):
                _validate_download(gguf_path)

    def test_empty_file(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = pathlib.Path(tmp_dir)
            gguf_path = tmp_path / "empty.gguf"
            gguf_path.write_text("")

            with pytest.raises(Exception, match="Download failed - file is empty"):
                _validate_download(gguf_path)


class TestDownloadWithRetry:
    """Test _download_with_retry function"""

    def test_download_success(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = pathlib.Path(tmp_dir)
            gguf_path = tmp_path / "model.gguf"

            with patch("subprocess.run") as mock_run:
                mock_run.return_value = None
                gguf_path.write_text("dummy content")

                result = _download_with_retry(
                    "https://example.com/model.gguf", gguf_path, max_retries=1
                )

                assert result is True
                mock_run.assert_called_once()

    def test_download_failure(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = pathlib.Path(tmp_dir)
            gguf_path = tmp_path / "model.gguf"

            with patch("subprocess.run", side_effect=subprocess.CalledProcessError(1, "curl")):
                result = _download_with_retry(
                    "https://example.com/model.gguf", gguf_path, max_retries=1
                )

                assert result is False


class TestGenerateReadmeContent:
    """Test _generate_readme_content function"""

    def test_readme_generation(self):
        content = _generate_readme_content("user/model", 3, 100)
        
        assert "base_model: user/model" in content
        assert "user/model" in content
        assert "3 split files" in content
        assert "100M each" in content

    def test_readme_generation_edge_cases(self):
        content = _generate_readme_content("", 1, 50)
        
        assert "base_model: " in content
        assert "1 split files" in content
        assert "50M each" in content


class TestSaveLocally:
    """Test save_locally function"""

    def test_save_success(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = pathlib.Path(tmp_dir)
            output_dir = tmp_path / "output"
            output_dir.mkdir()

            split_file1 = tmp_path / "model-00001-of-00002.gguf"
            split_file2 = tmp_path / "model-00002-of-00002.gguf"
            split_file1.write_text("dummy content 1")
            split_file2.write_text("dummy content 2")
            split_files = [split_file1, split_file2]

            result_path = save_locally(
                split_files, "user/model", "model.gguf", 100, output_dir
            )

            assert "model-sharded" in result_path
            assert (output_dir / "model-sharded").exists()
            assert (output_dir / "model-sharded" / "model-00001-of-00002.gguf").exists()
            assert (output_dir / "model-sharded" / "model-00002-of-00002.gguf").exists()
            assert (output_dir / "model-sharded" / "README.md").exists()


class TestUploadToHF:
    """Test upload_to_hf function"""

    def test_upload_success(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = pathlib.Path(tmp_dir)

            # Create dummy split files
            split_file1 = tmp_path / "model-00001-of-00002.gguf"
            split_file2 = tmp_path / "model-00002-of-00002.gguf"
            split_file1.write_text("dummy content 1")
            split_file2.write_text("dummy content 2")
            split_files = [split_file1, split_file2]

            oauth_token = Mock(spec=gr.OAuthToken)
            oauth_token.token = "test-token"

            with patch("src.processing.HfApi") as mock_api_class:
                mock_api = Mock()
                mock_api_class.return_value = mock_api

                result = upload_to_hf(
                    split_files, "test-repo", "user/model", oauth_token, 100, True
                )

                assert "https://huggingface.co/test-repo" in result
                assert mock_api.create_repo.called
                assert mock_api.upload_file.call_count == 3

    def test_upload_failure(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = pathlib.Path(tmp_dir)

            # Create dummy split files
            split_file1 = tmp_path / "model-00001-of-00002.gguf"
            split_file1.write_text("dummy content 1")
            split_files = [split_file1]

            oauth_token = Mock(spec=gr.OAuthToken)
            oauth_token.token = "test-token"

            with patch("src.processing.HfApi") as mock_api_class:
                mock_api = Mock()
                mock_api.create_repo.side_effect = Exception("Upload failed")
                mock_api_class.return_value = mock_api

                result = upload_to_hf(
                    split_files, "test-repo", "user/model", oauth_token, 100, True
                )

                assert result == ""


class TestCleanupTempFiles:
    """Test _cleanup_temp_files function"""

    def test_cleanup_success(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = pathlib.Path(tmp_dir)
            
            # Create dummy files to clean up
            gguf_path = tmp_path / "model.gguf"
            split_file1 = tmp_path / "model-00001-of-00002.gguf"
            split_file2 = tmp_path / "model-00002-of-00002.gguf"
            sharded_path = tmp_path / "sharded"
            sharded_path.mkdir()
            
            gguf_path.write_text("dummy content")
            split_file1.write_text("dummy content 1")
            split_file2.write_text("dummy content 2")
            
            assert gguf_path.exists()
            assert split_file1.exists()
            assert split_file2.exists()
            assert sharded_path.exists()
            
            _cleanup_temp_files(gguf_path, [split_file1, split_file2], sharded_path)
            
            assert not gguf_path.exists()
            assert not split_file1.exists()
            assert not split_file2.exists()
            assert not sharded_path.exists()

    def test_cleanup_with_permission_error(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = pathlib.Path(tmp_dir)
            
            # Create dummy files
            gguf_path = tmp_path / "model.gguf"
            split_file = tmp_path / "model-00001-of-00002.gguf"
            gguf_path.write_text("dummy content")
            split_file.write_text("dummy content")
            
            with patch.object(pathlib.Path, 'unlink', side_effect=PermissionError("Permission denied")):
                _cleanup_temp_files(gguf_path, [split_file])
                
                assert gguf_path.exists()
                assert split_file.exists()

    def test_cleanup_nonexistent_files(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = pathlib.Path(tmp_dir)
            
            # Create paths for nonexistent files
            gguf_path = tmp_path / "nonexistent.gguf"
            split_file = tmp_path / "nonexistent-split.gguf"
            
            _cleanup_temp_files(gguf_path, [split_file])
            
            assert not gguf_path.exists()
            assert not split_file.exists()


class TestSaveLocallyErrors:
    """Test error scenarios in save_locally function"""

    def test_save_with_invalid_output_dir(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = pathlib.Path(tmp_dir)
            
            # Create dummy split files
            split_file = tmp_path / "model-00001-of-00002.gguf"
            split_file.write_text("dummy content")
            split_files = [split_file]
            
            output_dir = tmp_path / "output"
            output_dir.mkdir()
            
            result = save_locally(
                split_files, "user/model", "model.gguf", 100, output_dir
            )
            
            assert result is not None
            assert "model-sharded" in result