gh-christianlouis-docuelevate/tests/test_extract_metadata_gpt.py

"""Comprehensive unit tests for app/tasks/extract_metadata_with_gpt.py module."""

import json
from unittest.mock import MagicMock, patch

import pytest

from app.tasks.extract_metadata_with_gpt import extract_json_from_text, extract_metadata_with_gpt


@pytest.mark.unit
class TestExtractJsonFromText:
    """Tests for extract_json_from_text function."""

    def test_extracts_json_from_backticks_with_json_tag(self):
        """Test extraction of JSON from triple-backtick block with json tag."""
        text = '```json\n{"key": "value", "num": 123}\n```'
        result = extract_json_from_text(text)
        assert result == '{"key": "value", "num": 123}'

    def test_extracts_json_from_backticks_no_lang(self):
        """Test extraction from backticks without language tag."""
        text = '```\n{"key": "value"}\n```'
        result = extract_json_from_text(text)
        assert result == '{"key": "value"}'

    def test_extracts_json_from_raw_text(self):
        """Test extraction from raw text with JSON."""
        text = 'Here is the result: {"key": "value", "nested": {"a": 1}} end.'
        result = extract_json_from_text(text)
        assert result == '{"key": "value", "nested": {"a": 1}}'

    def test_returns_none_for_no_json(self):
        """Test returns None when no JSON found."""
        text = "No JSON here at all, just plain text."
        result = extract_json_from_text(text)
        assert result is None

    def test_returns_none_for_incomplete_json(self):
        """Test returns None for incomplete JSON structures."""
        text = "Only opening brace: { but no closing"
        result = extract_json_from_text(text)
        assert result is None

    def test_extracts_complex_nested_json(self):
        """Test extraction of complex nested JSON."""
        text = '{"filename": "2024-01-01_Invoice", "tags": ["test", "invoice"], "metadata": {"amount": 100, "currency": "USD"}}'
        result = extract_json_from_text(text)
        parsed = json.loads(result)
        assert parsed["filename"] == "2024-01-01_Invoice"
        assert "tags" in parsed
        assert "metadata" in parsed
        assert parsed["metadata"]["amount"] == 100

    def test_extracts_first_json_when_multiple_present(self):
        """Test that extraction finds the outermost JSON object."""
        text = 'First: {"a": 1} and second: {"b": 2}'
        result = extract_json_from_text(text)
        # Should extract from first { to last }
        assert result is not None
        assert "{" in result and "}" in result


@pytest.mark.unit
class TestExtractMetadataWithGpt:
    """Tests for extract_metadata_with_gpt Celery task."""

    @patch("app.tasks.extract_metadata_with_gpt.embed_metadata_into_pdf")
    @patch("app.tasks.extract_metadata_with_gpt.log_task_progress")
    @patch("app.tasks.extract_metadata_with_gpt.get_ai_provider")
    def test_successful_metadata_extraction(self, mock_get_provider, mock_log_progress, mock_embed_task):
        """Test successful metadata extraction with valid AI provider response."""
        mock_provider = MagicMock()
        mock_provider.chat_completion.return_value = json.dumps(
            {
                "filename": "2024-01-15_Invoice_Amazon",
                "empfaenger": "John Doe",
                "absender": "Amazon",
                "correspondent": "Amazon",
                "kommunikationsart": "Rechnung",
                "kommunikationskategorie": "Finanz_und_Vertragsdokumente",
                "document_type": "Invoice",
                "tags": ["invoice", "amazon", "online-shopping"],
                "language": "de",
                "title": "Amazon Purchase Invoice",
                "confidence_score": 95,
                "reference_number": "INV-2024-001",
                "monetary_amounts": ["99.99 EUR"],
            }
        )
        mock_get_provider.return_value = mock_provider

        # Set task request context directly on the Celery task
        extract_metadata_with_gpt.request.id = "test-task-id"

        # Call the underlying function directly (not through Celery)
        result = extract_metadata_with_gpt.__wrapped__("test_invoice.pdf", "Invoice from Amazon for 99.99 EUR", 123)

        # Verify AI provider was called
        mock_provider.chat_completion.assert_called_once()
        call_kwargs = mock_provider.chat_completion.call_args[1]
        assert call_kwargs["temperature"] == 0
        assert len(call_kwargs["messages"]) == 2

        # Verify metadata was extracted correctly
        assert result["s3_file"] == "test_invoice.pdf"
        assert "metadata" in result
        assert result["metadata"]["document_type"] == "Invoice"
        assert result["metadata"]["correspondent"] == "Amazon"

        # Verify embed task was queued
        mock_embed_task.delay.assert_called_once()

        # Verify task progress was logged
        assert mock_log_progress.call_count >= 3

    @patch("app.tasks.extract_metadata_with_gpt.embed_metadata_into_pdf")
    @patch("app.tasks.extract_metadata_with_gpt.log_task_progress")
    @patch("app.tasks.extract_metadata_with_gpt.get_ai_provider")
    def test_handles_json_in_backticks(self, mock_get_provider, mock_log_progress, mock_embed_task):
        """Test extraction handles JSON wrapped in markdown code blocks."""
        mock_provider = MagicMock()
        mock_provider.chat_completion.return_value = (
            '```json\n{"filename": "test.pdf", "document_type": "Unknown"}\n```'
        )
        mock_get_provider.return_value = mock_provider

        extract_metadata_with_gpt.request.id = "test-task-id"

        result = extract_metadata_with_gpt.__wrapped__("test.pdf", "Sample text", 456)

        assert result["metadata"]["filename"] == "test.pdf"
        assert result["metadata"]["document_type"] == "Unknown"
        mock_embed_task.delay.assert_called_once()

    @patch("app.tasks.extract_metadata_with_gpt.embed_metadata_into_pdf")
    @patch("app.tasks.extract_metadata_with_gpt.log_task_progress")
    @patch("app.tasks.extract_metadata_with_gpt.get_ai_provider")
    def test_handles_invalid_json_response(self, mock_get_provider, mock_log_progress, mock_embed_task):
        """Test handling of invalid JSON in AI provider response."""
        mock_provider = MagicMock()
        mock_provider.chat_completion.return_value = "This is not valid JSON at all"
        mock_get_provider.return_value = mock_provider

        extract_metadata_with_gpt.request.id = "test-task-id"

        result = extract_metadata_with_gpt.__wrapped__("test.pdf", "Sample text", 789)

        assert result == {}
        mock_embed_task.delay.assert_not_called()
        # Verify failure was logged
        failure_calls = [call for call in mock_log_progress.call_args_list if "failure" in str(call)]
        assert len(failure_calls) > 0

    @patch("app.tasks.extract_metadata_with_gpt.embed_metadata_into_pdf")
    @patch("app.tasks.extract_metadata_with_gpt.log_task_progress")
    @patch("app.tasks.extract_metadata_with_gpt.get_ai_provider")
    def test_handles_api_exception(self, mock_get_provider, mock_log_progress, mock_embed_task):
        """Test handling of AI provider API exceptions."""
        mock_provider = MagicMock()
        mock_provider.chat_completion.side_effect = Exception("API Error: Rate limit exceeded")
        mock_get_provider.return_value = mock_provider

        extract_metadata_with_gpt.request.id = "test-task-id"

        result = extract_metadata_with_gpt.__wrapped__("test.pdf", "Sample text", 101)

        assert result == {}
        mock_embed_task.delay.assert_not_called()
        # Verify exception was logged
        failure_calls = [call for call in mock_log_progress.call_args_list if "failure" in str(call)]
        assert len(failure_calls) > 0

    @patch("app.tasks.extract_metadata_with_gpt.embed_metadata_into_pdf")
    @patch("app.tasks.extract_metadata_with_gpt.log_task_progress")
    @patch("app.tasks.extract_metadata_with_gpt.get_ai_provider")
    @patch("app.tasks.extract_metadata_with_gpt.SessionLocal")
    def test_retrieves_file_id_from_database_when_not_provided(
        self, mock_session_local, mock_get_provider, mock_log_progress, mock_embed_task
    ):
        """Test file_id retrieval from database when not provided."""
        mock_provider = MagicMock()
        mock_provider.chat_completion.return_value = '{"filename": "test.pdf", "document_type": "Unknown"}'
        mock_get_provider.return_value = mock_provider

        # Mock database session
        mock_db = MagicMock()
        mock_session_local.return_value.__enter__.return_value = mock_db
        mock_file_record = MagicMock()
        mock_file_record.id = 999
        mock_db.query.return_value.filter_by.return_value.first.return_value = mock_file_record

        # Mock file existence
        with patch("app.tasks.extract_metadata_with_gpt.os.path.exists", return_value=True):
            with patch("app.tasks.extract_metadata_with_gpt.settings.workdir", "/tmp"):
                extract_metadata_with_gpt.request.id = "test-task-id"

                result = extract_metadata_with_gpt.__wrapped__(
                    filename="test.pdf",
                    cleaned_text="Sample text",
                    file_id=None,  # Not provided
                )

                assert result["metadata"]["filename"] == "test.pdf"
                # Verify database was queried
                mock_db.query.assert_called_once()

    @patch("app.tasks.extract_metadata_with_gpt.embed_metadata_into_pdf")
    @patch("app.tasks.extract_metadata_with_gpt.log_task_progress")
    @patch("app.tasks.extract_metadata_with_gpt.get_ai_provider")
    def test_validates_filename_security(self, mock_get_provider, mock_log_progress, mock_embed_task):
        """Test filename validation to prevent path traversal."""
        mock_provider = MagicMock()
        # Try to inject a malicious filename
        mock_provider.chat_completion.return_value = json.dumps(
            {"filename": "../../../etc/passwd", "document_type": "Invoice"}
        )
        mock_get_provider.return_value = mock_provider

        extract_metadata_with_gpt.request.id = "test-task-id"

        result = extract_metadata_with_gpt.__wrapped__("test.pdf", "Sample text", 202)

        # Filename should be sanitized (empty or safe)
        assert result["metadata"]["filename"] == ""
        mock_embed_task.delay.assert_called_once()

    @patch("app.tasks.extract_metadata_with_gpt.embed_metadata_into_pdf")
    @patch("app.tasks.extract_metadata_with_gpt.log_task_progress")
    @patch("app.tasks.extract_metadata_with_gpt.get_ai_provider")
    def test_validates_filename_with_dots(self, mock_get_provider, mock_log_progress, mock_embed_task):
        """Test filename validation rejects '..' in filenames."""
        mock_provider = MagicMock()
        mock_provider.chat_completion.return_value = json.dumps(
            {"filename": "test..invoice.pdf", "document_type": "Invoice"}
        )
        mock_get_provider.return_value = mock_provider

        extract_metadata_with_gpt.request.id = "test-task-id"

        result = extract_metadata_with_gpt.__wrapped__("test.pdf", "Sample text", 303)

        # Filename with .. should be rejected
        assert result["metadata"]["filename"] == ""

    @patch("app.tasks.extract_metadata_with_gpt.embed_metadata_into_pdf")
    @patch("app.tasks.extract_metadata_with_gpt.log_task_progress")
    @patch("app.tasks.extract_metadata_with_gpt.get_ai_provider")
    def test_accepts_valid_filename(self, mock_get_provider, mock_log_progress, mock_embed_task):
        """Test that valid filenames are accepted."""
        mock_provider = MagicMock()
        mock_provider.chat_completion.return_value = json.dumps(
            {"filename": "2024-01-15_Invoice_Amazon.pdf", "document_type": "Invoice"}
        )
        mock_get_provider.return_value = mock_provider

        extract_metadata_with_gpt.request.id = "test-task-id"

        result = extract_metadata_with_gpt.__wrapped__("test.pdf", "Sample text", 404)

        # Valid filename should be preserved
        assert result["metadata"]["filename"] == "2024-01-15_Invoice_Amazon.pdf"

    @patch("app.tasks.extract_metadata_with_gpt.embed_metadata_into_pdf")
    @patch("app.tasks.extract_metadata_with_gpt.log_task_progress")
    @patch("app.tasks.extract_metadata_with_gpt.get_ai_provider")
    def test_handles_malformed_json_with_valid_structure(self, mock_get_provider, mock_log_progress, mock_embed_task):
        """Test handling of JSON that's parseable but missing expected fields."""
        mock_provider = MagicMock()
        mock_provider.chat_completion.return_value = '{"unexpected_field": "value"}'
        mock_get_provider.return_value = mock_provider

        extract_metadata_with_gpt.request.id = "test-task-id"

        result = extract_metadata_with_gpt.__wrapped__("test.pdf", "Sample text", 505)

        # Should still extract the JSON even if fields are unexpected
        assert "metadata" in result
        assert result["metadata"]["unexpected_field"] == "value"

    @patch("app.tasks.extract_metadata_with_gpt.embed_metadata_into_pdf")
    @patch("app.tasks.extract_metadata_with_gpt.log_task_progress")
    @patch("app.tasks.extract_metadata_with_gpt.get_ai_provider")
    def test_handles_absolute_path_filename(self, mock_get_provider, mock_log_progress, mock_embed_task):
        """Test handling when filename is provided as an absolute path."""
        mock_provider = MagicMock()
        mock_provider.chat_completion.return_value = '{"filename": "test.pdf", "document_type": "Unknown"}'
        mock_get_provider.return_value = mock_provider

        extract_metadata_with_gpt.request.id = "test-task-id"

        # Provide an absolute path as filename
        absolute_path = "/absolute/path/to/test.pdf"
        result = extract_metadata_with_gpt.__wrapped__(absolute_path, "Sample text", 606)

        # Should handle absolute path correctly
        assert result["s3_file"] == "test.pdf"  # Should extract basename
        assert "metadata" in result

    @patch("app.tasks.extract_metadata_with_gpt.embed_metadata_into_pdf")
    @patch("app.tasks.extract_metadata_with_gpt.log_task_progress")
    @patch("app.tasks.extract_metadata_with_gpt.get_ai_provider")
    @patch("app.tasks.extract_metadata_with_gpt.SessionLocal")
    def test_database_lookup_with_existing_file(
        self, mock_session_local, mock_get_provider, mock_log_progress, mock_embed_task
    ):
        """Test file_id retrieval when file exists on disk and in database."""
        mock_provider = MagicMock()
        mock_provider.chat_completion.return_value = '{"filename": "test.pdf", "document_type": "Unknown"}'
        mock_get_provider.return_value = mock_provider

        # Mock database session
        mock_db = MagicMock()
        mock_session_local.return_value.__enter__.return_value = mock_db
        mock_file_record = MagicMock()
        mock_file_record.id = 888
        mock_db.query.return_value.filter_by.return_value.first.return_value = mock_file_record

        # Mock file existence check
        with patch("app.tasks.extract_metadata_with_gpt.os.path.exists", return_value=True):
            with patch("app.tasks.extract_metadata_with_gpt.os.path.isabs", return_value=False):
                with patch("app.tasks.extract_metadata_with_gpt.settings.workdir", "/tmp"):
                    extract_metadata_with_gpt.request.id = "test-task-id"

                    result = extract_metadata_with_gpt.__wrapped__(
                        filename="test.pdf",
                        cleaned_text="Sample text",
                        file_id=None,  # Not provided, should look up
                    )

                    assert result["metadata"]["filename"] == "test.pdf"
                    # Verify database was queried
                    mock_db.query.assert_called_once()


@pytest.mark.unit
class TestModuleImports:
    """Tests for module import behaviour."""

    def test_module_imports_successfully(self):
        """Test that the module imports successfully."""
        import app.tasks.extract_metadata_with_gpt as mod

        assert mod is not None
        assert hasattr(mod, "extract_metadata_with_gpt")
        assert hasattr(mod, "extract_json_from_text")
        assert hasattr(mod, "get_ai_provider")