"""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