learningcircuit--local-deep-research/tests/research_library/services/test_pdf_storage_manager.py

"""Tests for PDF Storage Manager."""

from unittest.mock import MagicMock, patch


from local_deep_research.constants import (
    FILE_PATH_METADATA_ONLY,
)
from local_deep_research.research_library.services.pdf_storage_manager import (
    PDFStorageManager,
)


class TestPDFStorageManagerInit:
    """Tests for PDFStorageManager initialization."""

    def test_initializes_with_defaults(self, tmp_path):
        """Should initialize with default settings (3 GB)."""
        manager = PDFStorageManager(tmp_path, "database")
        assert manager.library_root == tmp_path.resolve()
        assert manager.storage_mode == "database"
        assert manager.max_pdf_size_bytes == 3072 * 1024 * 1024

    def test_initializes_with_custom_max_size(self, tmp_path):
        """Should use custom max PDF size."""
        manager = PDFStorageManager(tmp_path, "database", max_pdf_size_mb=50)
        assert manager.max_pdf_size_bytes == 50 * 1024 * 1024

    def test_accepts_valid_storage_modes(self, tmp_path):
        """Should accept all valid storage modes."""
        for mode in ("none", "filesystem", "database"):
            manager = PDFStorageManager(tmp_path, mode)
            assert manager.storage_mode == mode

    def test_defaults_to_none_for_invalid_mode(self, tmp_path):
        """Should default to 'none' for invalid storage mode."""
        manager = PDFStorageManager(tmp_path, "invalid_mode")
        assert manager.storage_mode == "none"


class TestPDFStorageManagerSavePdf:
    """Tests for save_pdf method."""

    def test_returns_none_for_none_mode(self, tmp_path, mock_pdf_content):
        """Should return None when storage mode is 'none'."""
        manager = PDFStorageManager(tmp_path, "none")
        mock_doc = MagicMock()
        mock_session = MagicMock()

        result, size = manager.save_pdf(
            mock_pdf_content, mock_doc, mock_session, "test.pdf"
        )

        assert result is None
        assert size == len(mock_pdf_content)

    def test_saves_to_database(self, tmp_path, mock_pdf_content):
        """Should save PDF to database when mode is 'database'."""
        manager = PDFStorageManager(tmp_path, "database")
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_session = MagicMock()
        mock_session.query.return_value.filter_by.return_value.first.return_value = None

        result, size = manager.save_pdf(
            mock_pdf_content, mock_doc, mock_session, "test.pdf"
        )

        assert result == "database"
        assert mock_doc.storage_mode == "database"
        assert mock_session.add.called

    def test_saves_to_filesystem(self, tmp_path, mock_pdf_content):
        """Should save PDF to filesystem when mode is 'filesystem'."""
        manager = PDFStorageManager(tmp_path, "filesystem")
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_session = MagicMock()

        # Mock the internal method to avoid file system operations
        with patch.object(
            manager,
            "_save_to_filesystem",
            return_value=tmp_path / "pdfs" / "test.pdf",
        ):
            result, size = manager.save_pdf(
                mock_pdf_content, mock_doc, mock_session, "test.pdf"
            )

        assert result is not None
        assert mock_doc.storage_mode == "filesystem"

    def test_rejects_oversized_pdf(self, tmp_path):
        """Should reject PDF that exceeds size limit."""
        manager = PDFStorageManager(tmp_path, "database", max_pdf_size_mb=1)
        mock_doc = MagicMock()
        mock_session = MagicMock()

        # Create content larger than 1MB
        large_content = b"x" * (2 * 1024 * 1024)

        result, size = manager.save_pdf(
            large_content, mock_doc, mock_session, "test.pdf"
        )

        assert result is None
        assert size == len(large_content)


class TestPDFStorageManagerLoadPdf:
    """Tests for load_pdf method."""

    def test_loads_from_database_first(self, tmp_path, mock_pdf_content):
        """Should try database first when loading."""
        manager = PDFStorageManager(tmp_path, "database")
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_session = MagicMock()

        # Mock database blob
        mock_blob = MagicMock()
        mock_blob.pdf_binary = mock_pdf_content
        mock_session.query.return_value.filter_by.return_value.first.return_value = mock_blob

        result = manager.load_pdf(mock_doc, mock_session)

        assert result == mock_pdf_content

    def test_falls_back_to_filesystem(self, tmp_path, mock_pdf_content):
        """Should fall back to filesystem if not in database."""
        manager = PDFStorageManager(tmp_path, "filesystem")
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_doc.file_path = "pdfs/test.pdf"
        mock_session = MagicMock()

        # No blob in database
        mock_session.query.return_value.filter_by.return_value.first.return_value = None

        # Create file on disk
        pdf_path = tmp_path / "pdfs"
        pdf_path.mkdir(parents=True, exist_ok=True)
        (pdf_path / "test.pdf").write_bytes(mock_pdf_content)

        result = manager.load_pdf(mock_doc, mock_session)

        assert result == mock_pdf_content

    def test_returns_none_when_not_found(self, tmp_path):
        """Should return None when PDF not found anywhere."""
        manager = PDFStorageManager(tmp_path, "database")
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_doc.file_path = None
        mock_session = MagicMock()

        mock_session.query.return_value.filter_by.return_value.first.return_value = None

        result = manager.load_pdf(mock_doc, mock_session)

        assert result is None


class TestPDFStorageManagerHasPdf:
    """Tests for has_pdf method."""

    def test_returns_false_for_non_pdf_file_type(self, tmp_path):
        """Should return False for non-PDF documents."""
        manager = PDFStorageManager(tmp_path, "database")
        mock_doc = MagicMock()
        mock_doc.file_type = "txt"
        mock_session = MagicMock()

        result = manager.has_pdf(mock_doc, mock_session)

        assert result is False

    def test_returns_false_for_non_pdf_type(self, tmp_path):
        """Should return False for non-PDF file types."""
        manager = PDFStorageManager(tmp_path, "database")
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_doc.file_type = "txt"  # Not a PDF
        mock_session = MagicMock()

        result = manager.has_pdf(mock_doc, mock_session)

        assert result is False

    def test_has_pdf_checks_database_first(self, tmp_path):
        """Should check database for blob existence."""
        manager = PDFStorageManager(tmp_path, "database")
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_doc.file_type = "pdf"
        mock_doc.file_path = None
        mock_session = MagicMock()

        # This test just verifies the method can be called
        # The actual implementation queries DocumentBlob
        try:
            manager.has_pdf(mock_doc, mock_session)
        except AttributeError:
            # Expected if DocumentBlob.id isn't accessible in test context
            pass


class TestPDFStorageManagerPdfExists:
    """Tests for pdf_exists classmethod."""

    def test_returns_false_for_non_pdf(self, tmp_path):
        """Should return False for non-PDF documents without requiring a storage mode."""
        mock_doc = MagicMock()
        mock_doc.file_type = "txt"
        mock_session = MagicMock()

        result = PDFStorageManager.pdf_exists(tmp_path, mock_doc, mock_session)

        assert result is False

    def test_returns_true_when_blob_exists(self, tmp_path):
        """Should find PDF in database via classmethod."""
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_doc.file_type = "pdf"
        mock_doc.file_path = None
        mock_session = MagicMock()

        # Simulate DocumentBlob found in database
        mock_session.query.return_value.filter_by.return_value.first.return_value = MagicMock()

        mock_blob_class = MagicMock()
        with patch(
            "local_deep_research.database.models.library.DocumentBlob",
            mock_blob_class,
        ):
            result = PDFStorageManager.pdf_exists(
                tmp_path, mock_doc, mock_session
            )

        assert result is True

    def test_returns_true_when_file_exists(self, tmp_path):
        """Should find PDF on filesystem via classmethod."""
        pdf_dir = tmp_path / "pdfs"
        pdf_dir.mkdir()
        (pdf_dir / "test.pdf").write_bytes(b"%PDF-1.4 test")

        mock_doc = MagicMock()
        mock_doc.id = "doc-456"
        mock_doc.file_type = "pdf"
        mock_doc.file_path = "pdfs/test.pdf"
        mock_session = MagicMock()

        # No blob in database
        mock_session.query.return_value.filter_by.return_value.first.return_value = None

        mock_blob_class = MagicMock()
        with patch(
            "local_deep_research.database.models.library.DocumentBlob",
            mock_blob_class,
        ):
            result = PDFStorageManager.pdf_exists(
                tmp_path, mock_doc, mock_session
            )

        assert result is True


class TestPDFStorageManagerDeletePdf:
    """Tests for delete_pdf method."""

    def test_deletes_database_blob(self, tmp_path):
        """Should delete blob from database."""
        manager = PDFStorageManager(tmp_path, "database")
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_doc.storage_mode = "database"
        mock_session = MagicMock()

        mock_blob = MagicMock()
        mock_session.query.return_value.filter_by.return_value.first.return_value = mock_blob

        result = manager.delete_pdf(mock_doc, mock_session)

        assert result is True
        mock_session.delete.assert_called_with(mock_blob)
        assert mock_doc.storage_mode == "none"

    def test_deletes_filesystem_file(self, tmp_path, mock_pdf_content):
        """Should delete file from filesystem."""
        manager = PDFStorageManager(tmp_path, "filesystem")
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_doc.storage_mode = "filesystem"
        mock_doc.file_path = "pdfs/test.pdf"
        mock_session = MagicMock()

        # Create file on disk
        pdf_path = tmp_path / "pdfs"
        pdf_path.mkdir(parents=True, exist_ok=True)
        file_path = pdf_path / "test.pdf"
        file_path.write_bytes(mock_pdf_content)

        result = manager.delete_pdf(mock_doc, mock_session)

        assert result is True
        assert not file_path.exists()
        assert mock_doc.storage_mode == "none"

    def test_returns_true_for_none_mode(self, tmp_path):
        """Should return True when nothing to delete."""
        manager = PDFStorageManager(tmp_path, "none")
        mock_doc = MagicMock()
        mock_doc.storage_mode = "none"
        mock_session = MagicMock()

        result = manager.delete_pdf(mock_doc, mock_session)

        assert result is True


class TestPDFStorageManagerUpgradeToPdf:
    """Tests for upgrade_to_pdf method."""

    def test_upgrades_text_only_document(self, tmp_path, mock_pdf_content):
        """Should add PDF to text-only document."""
        manager = PDFStorageManager(tmp_path, "database")
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_doc.storage_mode = "none"
        mock_session = MagicMock()

        mock_session.query.return_value.filter_by.return_value.first.return_value = None

        result = manager.upgrade_to_pdf(
            mock_doc, mock_pdf_content, mock_session
        )

        assert result is True
        assert mock_doc.storage_mode == "database"
        mock_session.add.assert_called()

    def test_skips_if_already_has_pdf(self, tmp_path, mock_pdf_content):
        """Should skip if document already has PDF storage."""
        manager = PDFStorageManager(tmp_path, "database")
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_doc.storage_mode = "database"
        mock_session = MagicMock()

        result = manager.upgrade_to_pdf(
            mock_doc, mock_pdf_content, mock_session
        )

        assert result is False

    def test_rejects_oversized_pdf(self, tmp_path):
        """Should reject PDF that exceeds size limit."""
        manager = PDFStorageManager(tmp_path, "database", max_pdf_size_mb=1)
        mock_doc = MagicMock()
        mock_doc.id = "doc-123"
        mock_doc.storage_mode = "none"
        mock_session = MagicMock()

        mock_session.query.return_value.filter_by.return_value.first.return_value = None

        large_content = b"x" * (2 * 1024 * 1024)

        result = manager.upgrade_to_pdf(mock_doc, large_content, mock_session)

        assert result is False


class TestPDFStorageManagerFilesystemWrite:
    """Tests for actual filesystem write operations."""

    def test_saves_to_filesystem_without_encoding_error(
        self, tmp_path, mock_pdf_content
    ):
        """Should save binary PDF to filesystem without encoding errors.

        Regression test for bug where write_file_verified passed encoding
        argument to binary mode open(), causing ValueError.
        """
        manager = PDFStorageManager(tmp_path, "filesystem")

        # Actually call _save_to_filesystem (not mocked)
        result_path = manager._save_to_filesystem(
            mock_pdf_content,
            "test_document.pdf",
            url="https://example.com/paper.pdf",
            resource_id=123,
        )

        # Verify file was created
        assert result_path.exists()
        assert result_path.suffix == ".pdf"

        # Verify content matches
        saved_content = result_path.read_bytes()
        assert saved_content == mock_pdf_content

    def test_filesystem_write_uses_settings_snapshot(
        self, tmp_path, mock_pdf_content
    ):
        """Should pass settings_snapshot to write_file_verified.

        Regression test for bug where settings couldn't be found because
        no settings_snapshot was passed to write_file_verified.
        """
        # Create manager with filesystem mode
        manager = PDFStorageManager(tmp_path, "filesystem")

        # This should work without raising FileWriteSecurityError
        # because the settings_snapshot is now passed correctly
        result_path = manager._save_to_filesystem(
            mock_pdf_content,
            "snapshot_test.pdf",
        )

        assert result_path.exists()


class TestPDFStorageManagerGenerateFilename:
    """Tests for _generate_filename method."""

    def test_generates_arxiv_filename(self, tmp_path):
        """Should generate proper filename for arXiv URLs."""
        manager = PDFStorageManager(tmp_path, "filesystem")
        filename = manager._generate_filename(
            "https://arxiv.org/pdf/2401.12345.pdf", None, "fallback.pdf"
        )
        assert "arxiv" in filename.lower()
        assert "2401.12345" in filename

    def test_generates_pmc_filename(self, tmp_path):
        """Should generate proper filename for PMC URLs."""
        manager = PDFStorageManager(tmp_path, "filesystem")
        filename = manager._generate_filename(
            "https://ncbi.nlm.nih.gov/pmc/articles/PMC1234567/pdf/",
            None,
            "fallback.pdf",
        )
        assert "pmc" in filename.lower() or "pubmed" in filename.lower()

    def test_uses_fallback_for_unknown_urls(self, tmp_path):
        """Should use fallback filename for unknown URLs."""
        manager = PDFStorageManager(tmp_path, "filesystem")
        filename = manager._generate_filename(
            "https://example.com/paper.pdf", None, "my_paper.pdf"
        )
        assert filename == "my_paper.pdf"


class TestPDFStorageManagerInferStorageMode:
    """Tests for _infer_storage_mode method."""

    def test_infers_database_from_blob(self, tmp_path):
        """Should infer database mode when blob exists."""
        manager = PDFStorageManager(tmp_path, "database")
        mock_doc = MagicMock()
        mock_doc.blob = MagicMock()
        mock_doc.file_path = None

        result = manager._infer_storage_mode(mock_doc)

        assert result == "database"

    def test_infers_filesystem_from_file_path(self, tmp_path):
        """Should infer filesystem mode when file_path exists."""
        manager = PDFStorageManager(tmp_path, "filesystem")
        mock_doc = MagicMock(spec=["file_path"])
        mock_doc.file_path = "pdfs/test.pdf"

        result = manager._infer_storage_mode(mock_doc)

        assert result == "filesystem"

    def test_infers_none_for_metadata_only(self, tmp_path):
        """Should infer none mode for metadata_only documents."""
        manager = PDFStorageManager(tmp_path, "database")
        mock_doc = MagicMock(spec=["file_path"])
        mock_doc.file_path = FILE_PATH_METADATA_ONLY

        result = manager._infer_storage_mode(mock_doc)

        assert result == "none"