829 lines
31 KiB
Python
829 lines
31 KiB
Python
"""Tests for storage helpers module."""
|
|
|
|
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from inference.web.services.storage_helpers import StorageHelper, get_storage_helper
|
|
from shared.storage import PREFIXES
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_storage() -> MagicMock:
|
|
"""Create a mock storage backend."""
|
|
storage = MagicMock()
|
|
storage.upload_bytes = MagicMock()
|
|
storage.download_bytes = MagicMock(return_value=b"test content")
|
|
storage.get_presigned_url = MagicMock(return_value="https://example.com/file")
|
|
storage.exists = MagicMock(return_value=True)
|
|
storage.delete = MagicMock(return_value=True)
|
|
storage.list_files = MagicMock(return_value=[])
|
|
return storage
|
|
|
|
|
|
@pytest.fixture
|
|
def helper(mock_storage: MagicMock) -> StorageHelper:
|
|
"""Create a storage helper with mock backend."""
|
|
return StorageHelper(storage=mock_storage)
|
|
|
|
|
|
class TestStorageHelperInit:
|
|
"""Tests for StorageHelper initialization."""
|
|
|
|
def test_init_with_storage(self, mock_storage: MagicMock) -> None:
|
|
"""Should use provided storage backend."""
|
|
helper = StorageHelper(storage=mock_storage)
|
|
assert helper.storage is mock_storage
|
|
|
|
def test_storage_property(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Storage property should return the backend."""
|
|
assert helper.storage is mock_storage
|
|
|
|
|
|
class TestDocumentOperations:
|
|
"""Tests for document storage operations."""
|
|
|
|
def test_upload_document(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should upload document with correct path."""
|
|
doc_id, path = helper.upload_document(b"pdf content", "invoice.pdf", "doc123")
|
|
|
|
assert doc_id == "doc123"
|
|
assert path == "documents/doc123.pdf"
|
|
mock_storage.upload_bytes.assert_called_once_with(
|
|
b"pdf content", "documents/doc123.pdf", overwrite=True
|
|
)
|
|
|
|
def test_upload_document_generates_id(self, helper: StorageHelper) -> None:
|
|
"""Should generate document ID if not provided."""
|
|
doc_id, path = helper.upload_document(b"content", "file.pdf")
|
|
|
|
assert doc_id is not None
|
|
assert len(doc_id) > 0
|
|
assert path.startswith("documents/")
|
|
|
|
def test_download_document(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should download document from correct path."""
|
|
content = helper.download_document("doc123")
|
|
|
|
assert content == b"test content"
|
|
mock_storage.download_bytes.assert_called_once_with("documents/doc123.pdf")
|
|
|
|
def test_get_document_url(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should get presigned URL for document."""
|
|
url = helper.get_document_url("doc123", expires_in_seconds=7200)
|
|
|
|
assert url == "https://example.com/file"
|
|
mock_storage.get_presigned_url.assert_called_once_with(
|
|
"documents/doc123.pdf", 7200
|
|
)
|
|
|
|
def test_document_exists(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should check document existence."""
|
|
exists = helper.document_exists("doc123")
|
|
|
|
assert exists is True
|
|
mock_storage.exists.assert_called_once_with("documents/doc123.pdf")
|
|
|
|
def test_delete_document(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should delete document."""
|
|
result = helper.delete_document("doc123")
|
|
|
|
assert result is True
|
|
mock_storage.delete.assert_called_once_with("documents/doc123.pdf")
|
|
|
|
|
|
class TestImageOperations:
|
|
"""Tests for image storage operations."""
|
|
|
|
def test_save_page_image(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should save page image with correct path."""
|
|
path = helper.save_page_image("doc123", 1, b"image data")
|
|
|
|
assert path == "images/doc123/page_1.png"
|
|
mock_storage.upload_bytes.assert_called_once_with(
|
|
b"image data", "images/doc123/page_1.png", overwrite=True
|
|
)
|
|
|
|
def test_get_page_image(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should get page image from correct path."""
|
|
content = helper.get_page_image("doc123", 2)
|
|
|
|
assert content == b"test content"
|
|
mock_storage.download_bytes.assert_called_once_with("images/doc123/page_2.png")
|
|
|
|
def test_get_page_image_url(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should get presigned URL for page image."""
|
|
url = helper.get_page_image_url("doc123", 3)
|
|
|
|
assert url == "https://example.com/file"
|
|
mock_storage.get_presigned_url.assert_called_once_with(
|
|
"images/doc123/page_3.png", 3600
|
|
)
|
|
|
|
def test_delete_document_images(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should delete all images for a document."""
|
|
mock_storage.list_files.return_value = [
|
|
"images/doc123/page_1.png",
|
|
"images/doc123/page_2.png",
|
|
]
|
|
|
|
deleted = helper.delete_document_images("doc123")
|
|
|
|
assert deleted == 2
|
|
mock_storage.list_files.assert_called_once_with("images/doc123/")
|
|
|
|
def test_list_document_images(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should list all images for a document."""
|
|
mock_storage.list_files.return_value = ["images/doc123/page_1.png"]
|
|
|
|
images = helper.list_document_images("doc123")
|
|
|
|
assert images == ["images/doc123/page_1.png"]
|
|
|
|
|
|
class TestUploadOperations:
|
|
"""Tests for upload staging operations."""
|
|
|
|
def test_save_upload(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should save upload to correct path."""
|
|
path = helper.save_upload(b"content", "file.pdf")
|
|
|
|
assert path == "uploads/file.pdf"
|
|
mock_storage.upload_bytes.assert_called_once()
|
|
|
|
def test_save_upload_with_subfolder(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should save upload with subfolder."""
|
|
path = helper.save_upload(b"content", "file.pdf", "async")
|
|
|
|
assert path == "uploads/async/file.pdf"
|
|
|
|
def test_get_upload(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should get upload from correct path."""
|
|
content = helper.get_upload("file.pdf", "async")
|
|
|
|
mock_storage.download_bytes.assert_called_once_with("uploads/async/file.pdf")
|
|
|
|
def test_delete_upload(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should delete upload."""
|
|
result = helper.delete_upload("file.pdf")
|
|
|
|
assert result is True
|
|
mock_storage.delete.assert_called_once_with("uploads/file.pdf")
|
|
|
|
|
|
class TestResultOperations:
|
|
"""Tests for result file operations."""
|
|
|
|
def test_save_result(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should save result to correct path."""
|
|
path = helper.save_result(b"result data", "output.json")
|
|
|
|
assert path == "results/output.json"
|
|
mock_storage.upload_bytes.assert_called_once()
|
|
|
|
def test_get_result(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should get result from correct path."""
|
|
content = helper.get_result("output.json")
|
|
|
|
mock_storage.download_bytes.assert_called_once_with("results/output.json")
|
|
|
|
def test_get_result_url(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should get presigned URL for result."""
|
|
url = helper.get_result_url("output.json")
|
|
|
|
mock_storage.get_presigned_url.assert_called_once_with("results/output.json", 3600)
|
|
|
|
def test_result_exists(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should check result existence."""
|
|
exists = helper.result_exists("output.json")
|
|
|
|
assert exists is True
|
|
mock_storage.exists.assert_called_once_with("results/output.json")
|
|
|
|
|
|
class TestExportOperations:
|
|
"""Tests for export file operations."""
|
|
|
|
def test_save_export(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should save export to correct path."""
|
|
path = helper.save_export(b"export data", "exp123", "dataset.zip")
|
|
|
|
assert path == "exports/exp123/dataset.zip"
|
|
mock_storage.upload_bytes.assert_called_once()
|
|
|
|
def test_get_export_url(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should get presigned URL for export."""
|
|
url = helper.get_export_url("exp123", "dataset.zip")
|
|
|
|
mock_storage.get_presigned_url.assert_called_once_with(
|
|
"exports/exp123/dataset.zip", 3600
|
|
)
|
|
|
|
|
|
class TestRawPdfOperations:
|
|
"""Tests for raw PDF operations (legacy compatibility)."""
|
|
|
|
def test_save_raw_pdf(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should save raw PDF to correct path."""
|
|
path = helper.save_raw_pdf(b"pdf data", "invoice.pdf")
|
|
|
|
assert path == "raw_pdfs/invoice.pdf"
|
|
mock_storage.upload_bytes.assert_called_once()
|
|
|
|
def test_get_raw_pdf(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should get raw PDF from correct path."""
|
|
content = helper.get_raw_pdf("invoice.pdf")
|
|
|
|
mock_storage.download_bytes.assert_called_once_with("raw_pdfs/invoice.pdf")
|
|
|
|
def test_raw_pdf_exists(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should check raw PDF existence."""
|
|
exists = helper.raw_pdf_exists("invoice.pdf")
|
|
|
|
assert exists is True
|
|
mock_storage.exists.assert_called_once_with("raw_pdfs/invoice.pdf")
|
|
|
|
|
|
class TestAdminImageOperations:
|
|
"""Tests for admin image storage operations."""
|
|
|
|
def test_save_admin_image(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should save admin image with correct path."""
|
|
path = helper.save_admin_image("doc123", 1, b"image data")
|
|
|
|
assert path == "admin_images/doc123/page_1.png"
|
|
mock_storage.upload_bytes.assert_called_once_with(
|
|
b"image data", "admin_images/doc123/page_1.png", overwrite=True
|
|
)
|
|
|
|
def test_get_admin_image(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should get admin image from correct path."""
|
|
content = helper.get_admin_image("doc123", 2)
|
|
|
|
assert content == b"test content"
|
|
mock_storage.download_bytes.assert_called_once_with("admin_images/doc123/page_2.png")
|
|
|
|
def test_get_admin_image_url(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should get presigned URL for admin image."""
|
|
url = helper.get_admin_image_url("doc123", 3)
|
|
|
|
assert url == "https://example.com/file"
|
|
mock_storage.get_presigned_url.assert_called_once_with(
|
|
"admin_images/doc123/page_3.png", 3600
|
|
)
|
|
|
|
def test_admin_image_exists(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should check admin image existence."""
|
|
exists = helper.admin_image_exists("doc123", 1)
|
|
|
|
assert exists is True
|
|
mock_storage.exists.assert_called_once_with("admin_images/doc123/page_1.png")
|
|
|
|
def test_get_admin_image_path(self, helper: StorageHelper) -> None:
|
|
"""Should return correct admin image path."""
|
|
path = helper.get_admin_image_path("doc123", 2)
|
|
|
|
assert path == "admin_images/doc123/page_2.png"
|
|
|
|
def test_list_admin_images(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should list all admin images for a document."""
|
|
mock_storage.list_files.return_value = [
|
|
"admin_images/doc123/page_1.png",
|
|
"admin_images/doc123/page_2.png",
|
|
]
|
|
|
|
images = helper.list_admin_images("doc123")
|
|
|
|
assert images == ["admin_images/doc123/page_1.png", "admin_images/doc123/page_2.png"]
|
|
mock_storage.list_files.assert_called_once_with("admin_images/doc123/")
|
|
|
|
def test_delete_admin_images(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should delete all admin images for a document."""
|
|
mock_storage.list_files.return_value = [
|
|
"admin_images/doc123/page_1.png",
|
|
"admin_images/doc123/page_2.png",
|
|
]
|
|
|
|
deleted = helper.delete_admin_images("doc123")
|
|
|
|
assert deleted == 2
|
|
mock_storage.list_files.assert_called_once_with("admin_images/doc123/")
|
|
|
|
|
|
class TestGetLocalPath:
|
|
"""Tests for get_local_path method."""
|
|
|
|
def test_get_admin_image_local_path_with_local_storage(self) -> None:
|
|
"""Should return local path when using local storage backend."""
|
|
from pathlib import Path
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
# Create a test image
|
|
test_path = Path(temp_dir) / "admin_images" / "doc123"
|
|
test_path.mkdir(parents=True, exist_ok=True)
|
|
(test_path / "page_1.png").write_bytes(b"test image")
|
|
|
|
local_path = helper.get_admin_image_local_path("doc123", 1)
|
|
|
|
assert local_path is not None
|
|
assert local_path.exists()
|
|
assert local_path.name == "page_1.png"
|
|
|
|
def test_get_admin_image_local_path_returns_none_for_cloud(
|
|
self, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should return None when storage doesn't support local paths."""
|
|
# Mock storage without get_local_path method (simulating cloud storage)
|
|
mock_storage.get_local_path = MagicMock(return_value=None)
|
|
helper = StorageHelper(storage=mock_storage)
|
|
|
|
local_path = helper.get_admin_image_local_path("doc123", 1)
|
|
|
|
assert local_path is None
|
|
|
|
def test_get_admin_image_local_path_nonexistent_file(self) -> None:
|
|
"""Should return None when file doesn't exist."""
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
local_path = helper.get_admin_image_local_path("nonexistent", 1)
|
|
|
|
assert local_path is None
|
|
|
|
|
|
class TestGetAdminImageDimensions:
|
|
"""Tests for get_admin_image_dimensions method."""
|
|
|
|
def test_get_dimensions_with_local_storage(self) -> None:
|
|
"""Should return image dimensions when using local storage."""
|
|
from pathlib import Path
|
|
from shared.storage.local import LocalStorageBackend
|
|
from PIL import Image
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
# Create a test image with known dimensions
|
|
test_path = Path(temp_dir) / "admin_images" / "doc123"
|
|
test_path.mkdir(parents=True, exist_ok=True)
|
|
img = Image.new("RGB", (800, 600), color="white")
|
|
img.save(test_path / "page_1.png")
|
|
|
|
dimensions = helper.get_admin_image_dimensions("doc123", 1)
|
|
|
|
assert dimensions == (800, 600)
|
|
|
|
def test_get_dimensions_nonexistent_file(self) -> None:
|
|
"""Should return None when file doesn't exist."""
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
dimensions = helper.get_admin_image_dimensions("nonexistent", 1)
|
|
|
|
assert dimensions is None
|
|
|
|
|
|
class TestGetStorageHelper:
|
|
"""Tests for get_storage_helper function."""
|
|
|
|
def test_returns_helper_instance(self) -> None:
|
|
"""Should return a StorageHelper instance."""
|
|
with patch("inference.web.services.storage_helpers.get_default_storage") as mock_get:
|
|
mock_get.return_value = MagicMock()
|
|
# Reset the global helper
|
|
import inference.web.services.storage_helpers as module
|
|
module._default_helper = None
|
|
|
|
helper = get_storage_helper()
|
|
|
|
assert isinstance(helper, StorageHelper)
|
|
|
|
def test_returns_same_instance(self) -> None:
|
|
"""Should return the same instance on subsequent calls."""
|
|
with patch("inference.web.services.storage_helpers.get_default_storage") as mock_get:
|
|
mock_get.return_value = MagicMock()
|
|
import inference.web.services.storage_helpers as module
|
|
module._default_helper = None
|
|
|
|
helper1 = get_storage_helper()
|
|
helper2 = get_storage_helper()
|
|
|
|
assert helper1 is helper2
|
|
|
|
|
|
class TestDeleteResult:
|
|
"""Tests for delete_result method."""
|
|
|
|
def test_delete_result(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should delete result file."""
|
|
result = helper.delete_result("output.json")
|
|
|
|
assert result is True
|
|
mock_storage.delete.assert_called_once_with("results/output.json")
|
|
|
|
|
|
class TestResultLocalPath:
|
|
"""Tests for get_result_local_path method."""
|
|
|
|
def test_get_result_local_path_with_local_storage(self) -> None:
|
|
"""Should return local path when using local storage backend."""
|
|
from pathlib import Path
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
# Create a test result file
|
|
test_path = Path(temp_dir) / "results"
|
|
test_path.mkdir(parents=True, exist_ok=True)
|
|
(test_path / "output.json").write_bytes(b"test result")
|
|
|
|
local_path = helper.get_result_local_path("output.json")
|
|
|
|
assert local_path is not None
|
|
assert local_path.exists()
|
|
assert local_path.name == "output.json"
|
|
|
|
def test_get_result_local_path_returns_none_for_cloud(
|
|
self, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should return None when storage doesn't support local paths."""
|
|
helper = StorageHelper(storage=mock_storage)
|
|
local_path = helper.get_result_local_path("output.json")
|
|
assert local_path is None
|
|
|
|
def test_get_result_local_path_nonexistent_file(self) -> None:
|
|
"""Should return None when file doesn't exist."""
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
local_path = helper.get_result_local_path("nonexistent.json")
|
|
|
|
assert local_path is None
|
|
|
|
|
|
class TestResultsBasePath:
|
|
"""Tests for get_results_base_path method."""
|
|
|
|
def test_get_results_base_path_with_local_storage(self) -> None:
|
|
"""Should return base path when using local storage."""
|
|
from pathlib import Path
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
base_path = helper.get_results_base_path()
|
|
|
|
assert base_path is not None
|
|
assert base_path.exists()
|
|
assert base_path.name == "results"
|
|
|
|
def test_get_results_base_path_returns_none_for_cloud(
|
|
self, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should return None when not using local storage."""
|
|
helper = StorageHelper(storage=mock_storage)
|
|
base_path = helper.get_results_base_path()
|
|
assert base_path is None
|
|
|
|
|
|
class TestUploadLocalPath:
|
|
"""Tests for get_upload_local_path method."""
|
|
|
|
def test_get_upload_local_path_with_local_storage(self) -> None:
|
|
"""Should return local path when using local storage backend."""
|
|
from pathlib import Path
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
# Create a test upload file
|
|
test_path = Path(temp_dir) / "uploads"
|
|
test_path.mkdir(parents=True, exist_ok=True)
|
|
(test_path / "file.pdf").write_bytes(b"test upload")
|
|
|
|
local_path = helper.get_upload_local_path("file.pdf")
|
|
|
|
assert local_path is not None
|
|
assert local_path.exists()
|
|
assert local_path.name == "file.pdf"
|
|
|
|
def test_get_upload_local_path_with_subfolder(self) -> None:
|
|
"""Should return local path with subfolder."""
|
|
from pathlib import Path
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
# Create a test upload file with subfolder
|
|
test_path = Path(temp_dir) / "uploads" / "async"
|
|
test_path.mkdir(parents=True, exist_ok=True)
|
|
(test_path / "file.pdf").write_bytes(b"test upload")
|
|
|
|
local_path = helper.get_upload_local_path("file.pdf", "async")
|
|
|
|
assert local_path is not None
|
|
assert local_path.exists()
|
|
|
|
def test_get_upload_local_path_returns_none_for_cloud(
|
|
self, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should return None when not using local storage."""
|
|
helper = StorageHelper(storage=mock_storage)
|
|
local_path = helper.get_upload_local_path("file.pdf")
|
|
assert local_path is None
|
|
|
|
|
|
class TestUploadsBasePath:
|
|
"""Tests for get_uploads_base_path method."""
|
|
|
|
def test_get_uploads_base_path_with_local_storage(self) -> None:
|
|
"""Should return base path when using local storage."""
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
base_path = helper.get_uploads_base_path()
|
|
|
|
assert base_path is not None
|
|
assert base_path.exists()
|
|
assert base_path.name == "uploads"
|
|
|
|
def test_get_uploads_base_path_with_subfolder(self) -> None:
|
|
"""Should return base path with subfolder."""
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
base_path = helper.get_uploads_base_path("async")
|
|
|
|
assert base_path is not None
|
|
assert base_path.exists()
|
|
assert base_path.name == "async"
|
|
|
|
def test_get_uploads_base_path_returns_none_for_cloud(
|
|
self, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should return None when not using local storage."""
|
|
helper = StorageHelper(storage=mock_storage)
|
|
base_path = helper.get_uploads_base_path()
|
|
assert base_path is None
|
|
|
|
|
|
class TestUploadExists:
|
|
"""Tests for upload_exists method."""
|
|
|
|
def test_upload_exists(self, helper: StorageHelper, mock_storage: MagicMock) -> None:
|
|
"""Should check upload existence."""
|
|
exists = helper.upload_exists("file.pdf")
|
|
|
|
assert exists is True
|
|
mock_storage.exists.assert_called_once_with("uploads/file.pdf")
|
|
|
|
def test_upload_exists_with_subfolder(
|
|
self, helper: StorageHelper, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should check upload existence with subfolder."""
|
|
helper.upload_exists("file.pdf", "async")
|
|
|
|
mock_storage.exists.assert_called_once_with("uploads/async/file.pdf")
|
|
|
|
|
|
class TestDatasetsBasePath:
|
|
"""Tests for get_datasets_base_path method."""
|
|
|
|
def test_get_datasets_base_path_with_local_storage(self) -> None:
|
|
"""Should return base path when using local storage."""
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
base_path = helper.get_datasets_base_path()
|
|
|
|
assert base_path is not None
|
|
assert base_path.exists()
|
|
assert base_path.name == "datasets"
|
|
|
|
def test_get_datasets_base_path_returns_none_for_cloud(
|
|
self, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should return None when not using local storage."""
|
|
helper = StorageHelper(storage=mock_storage)
|
|
base_path = helper.get_datasets_base_path()
|
|
assert base_path is None
|
|
|
|
|
|
class TestAdminImagesBasePath:
|
|
"""Tests for get_admin_images_base_path method."""
|
|
|
|
def test_get_admin_images_base_path_with_local_storage(self) -> None:
|
|
"""Should return base path when using local storage."""
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
base_path = helper.get_admin_images_base_path()
|
|
|
|
assert base_path is not None
|
|
assert base_path.exists()
|
|
assert base_path.name == "admin_images"
|
|
|
|
def test_get_admin_images_base_path_returns_none_for_cloud(
|
|
self, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should return None when not using local storage."""
|
|
helper = StorageHelper(storage=mock_storage)
|
|
base_path = helper.get_admin_images_base_path()
|
|
assert base_path is None
|
|
|
|
|
|
class TestRawPdfsBasePath:
|
|
"""Tests for get_raw_pdfs_base_path method."""
|
|
|
|
def test_get_raw_pdfs_base_path_with_local_storage(self) -> None:
|
|
"""Should return base path when using local storage."""
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
base_path = helper.get_raw_pdfs_base_path()
|
|
|
|
assert base_path is not None
|
|
assert base_path.exists()
|
|
assert base_path.name == "raw_pdfs"
|
|
|
|
def test_get_raw_pdfs_base_path_returns_none_for_cloud(
|
|
self, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should return None when not using local storage."""
|
|
helper = StorageHelper(storage=mock_storage)
|
|
base_path = helper.get_raw_pdfs_base_path()
|
|
assert base_path is None
|
|
|
|
|
|
class TestRawPdfLocalPath:
|
|
"""Tests for get_raw_pdf_local_path method."""
|
|
|
|
def test_get_raw_pdf_local_path_with_local_storage(self) -> None:
|
|
"""Should return local path when using local storage backend."""
|
|
from pathlib import Path
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
# Create a test raw PDF
|
|
test_path = Path(temp_dir) / "raw_pdfs"
|
|
test_path.mkdir(parents=True, exist_ok=True)
|
|
(test_path / "invoice.pdf").write_bytes(b"test pdf")
|
|
|
|
local_path = helper.get_raw_pdf_local_path("invoice.pdf")
|
|
|
|
assert local_path is not None
|
|
assert local_path.exists()
|
|
assert local_path.name == "invoice.pdf"
|
|
|
|
def test_get_raw_pdf_local_path_returns_none_for_cloud(
|
|
self, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should return None when not using local storage."""
|
|
helper = StorageHelper(storage=mock_storage)
|
|
local_path = helper.get_raw_pdf_local_path("invoice.pdf")
|
|
assert local_path is None
|
|
|
|
|
|
class TestRawPdfPath:
|
|
"""Tests for get_raw_pdf_path method."""
|
|
|
|
def test_get_raw_pdf_path(self, helper: StorageHelper) -> None:
|
|
"""Should return correct storage path."""
|
|
path = helper.get_raw_pdf_path("invoice.pdf")
|
|
assert path == "raw_pdfs/invoice.pdf"
|
|
|
|
|
|
class TestAutolabelOutputPath:
|
|
"""Tests for get_autolabel_output_path method."""
|
|
|
|
def test_get_autolabel_output_path_with_local_storage(self) -> None:
|
|
"""Should return output path when using local storage."""
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
output_path = helper.get_autolabel_output_path()
|
|
|
|
assert output_path is not None
|
|
assert output_path.exists()
|
|
assert output_path.name == "autolabel_output"
|
|
|
|
def test_get_autolabel_output_path_returns_none_for_cloud(
|
|
self, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should return None when not using local storage."""
|
|
helper = StorageHelper(storage=mock_storage)
|
|
output_path = helper.get_autolabel_output_path()
|
|
assert output_path is None
|
|
|
|
|
|
class TestTrainingDataPath:
|
|
"""Tests for get_training_data_path method."""
|
|
|
|
def test_get_training_data_path_with_local_storage(self) -> None:
|
|
"""Should return training path when using local storage."""
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
training_path = helper.get_training_data_path()
|
|
|
|
assert training_path is not None
|
|
assert training_path.exists()
|
|
assert training_path.name == "training"
|
|
|
|
def test_get_training_data_path_returns_none_for_cloud(
|
|
self, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should return None when not using local storage."""
|
|
helper = StorageHelper(storage=mock_storage)
|
|
training_path = helper.get_training_data_path()
|
|
assert training_path is None
|
|
|
|
|
|
class TestExportsBasePath:
|
|
"""Tests for get_exports_base_path method."""
|
|
|
|
def test_get_exports_base_path_with_local_storage(self) -> None:
|
|
"""Should return base path when using local storage."""
|
|
from shared.storage.local import LocalStorageBackend
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
storage = LocalStorageBackend(temp_dir)
|
|
helper = StorageHelper(storage=storage)
|
|
|
|
base_path = helper.get_exports_base_path()
|
|
|
|
assert base_path is not None
|
|
assert base_path.exists()
|
|
assert base_path.name == "exports"
|
|
|
|
def test_get_exports_base_path_returns_none_for_cloud(
|
|
self, mock_storage: MagicMock
|
|
) -> None:
|
|
"""Should return None when not using local storage."""
|
|
helper = StorageHelper(storage=mock_storage)
|
|
base_path = helper.get_exports_base_path()
|
|
assert base_path is None
|