python/tests/test_code_extraction_source_id.py

"""
Test that code extraction uses the correct source_id.

This test ensures that the fix for using hash-based source_ids
instead of domain-based source_ids works correctly.
"""

from unittest.mock import AsyncMock, Mock

import pytest

from src.server.services.crawling.code_extraction_service import CodeExtractionService
from src.server.services.storage.document_storage import DocumentStorageFacade


class TestCodeExtractionSourceId:
    """Test that code extraction properly uses the provided source_id."""

    @pytest.mark.asyncio
    async def test_code_extraction_uses_provided_source_id(self):
        """Test that code extraction uses the hash-based source_id, not domain."""
        # Create mock supabase client
        mock_supabase = Mock()
        mock_supabase.table.return_value.select.return_value.eq.return_value.execute.return_value.data = []

        # Create service instance
        code_service = CodeExtractionService(mock_supabase)

        # Track what gets passed to the internal extraction method
        extracted_blocks = []

        async def mock_extract_blocks(
            crawl_results, source_id, progress_callback=None, start=0, end=100, cancellation_check=None
        ):
            # Simulate finding code blocks and verify source_id is passed correctly
            for doc in crawl_results:
                extracted_blocks.append(
                    {
                        "block": {"code": "print('hello')", "language": "python"},
                        "source_url": doc["url"],
                        "source_id": source_id,  # This should be the provided source_id
                    }
                )
            return extracted_blocks

        code_service._extract_code_blocks_from_documents = mock_extract_blocks
        code_service._generate_code_summaries = AsyncMock(return_value=[{"summary": "Test code"}])
        code_service._prepare_code_examples_for_storage = Mock(
            return_value=[{"source_id": extracted_blocks[0]["source_id"] if extracted_blocks else None}]
        )
        code_service._store_code_examples = AsyncMock(return_value=1)

        # Test data
        crawl_results = [{"url": "https://docs.mem0.ai/example", "markdown": "```python\nprint('hello')\n```"}]

        url_to_full_document = {"https://docs.mem0.ai/example": "Full content with code"}

        # The correct hash-based source_id
        correct_source_id = "393224e227ba92eb"

        # Call the method with the correct source_id
        await code_service.extract_and_store_code_examples(crawl_results, url_to_full_document, correct_source_id, None)

        # Verify that extracted blocks use the correct source_id
        assert len(extracted_blocks) > 0, "Should have extracted at least one code block"

        for block in extracted_blocks:
            # Check that it's using the hash-based source_id, not the domain
            assert block["source_id"] == correct_source_id, (
                f"Should use hash-based source_id '{correct_source_id}', not domain"
            )
            assert block["source_id"] != "docs.mem0.ai", "Should NOT use domain-based source_id"

    @pytest.mark.asyncio
    async def test_document_storage_passes_source_id(self):
        """Test that DocumentStorageFacade passes source_id to code extraction."""
        # Create mock supabase client
        mock_supabase = Mock()

        # Create DocumentStorageFacade instance
        doc_storage = DocumentStorageFacade(mock_supabase)

        # Mock the code extraction service
        mock_extract = AsyncMock(return_value=5)
        doc_storage.code_extraction_service.extract_and_store_code_examples = mock_extract

        # Test data
        crawl_results = [{"url": "https://example.com", "markdown": "test"}]
        url_to_full_document = {"https://example.com": "test content"}
        source_id = "abc123def456"

        # Call the wrapper method
        result = await doc_storage.store_code_examples(crawl_results, url_to_full_document, source_id, None)

        # Verify the correct source_id was passed (now with cancellation_check parameter)
        mock_extract.assert_called_once()
        args, kwargs = mock_extract.call_args
        assert args[0] == crawl_results
        assert args[1] == url_to_full_document
        assert args[2] == source_id
        assert args[3] is None  # progress_callback
        assert args[4] == 85  # start_progress (default value)
        assert args[5] == 95  # end_progress (default value)
        assert args[6] is None  # cancellation_check
        assert result == 5

    @pytest.mark.asyncio
    async def test_no_domain_extraction_from_url(self):
        """Test that we're NOT extracting domain from URL anymore."""
        mock_supabase = Mock()
        mock_supabase.table.return_value.select.return_value.eq.return_value.execute.return_value.data = []

        code_service = CodeExtractionService(mock_supabase)

        # Patch internal methods
        code_service._get_setting = AsyncMock(return_value=True)

        # Create a mock that will track what source_id is used
        source_ids_seen = []

        async def track_source_id(
            crawl_results, source_id, progress_callback=None, start=0, end=100, cancellation_check=None
        ):
            source_ids_seen.append(source_id)
            return []  # Return empty list to skip further processing

        code_service._extract_code_blocks_from_documents = track_source_id

        # Test with various URLs that would produce different domains
        test_cases = [
            ("https://github.com/example/repo", "github123abc"),
            ("https://docs.python.org/guide", "python456def"),
            ("https://api.openai.com/v1", "openai789ghi"),
        ]

        for url, expected_source_id in test_cases:
            source_ids_seen.clear()

            crawl_results = [{"url": url, "markdown": "# Test"}]
            url_to_full_document = {url: "Full content"}

            await code_service.extract_and_store_code_examples(
                crawl_results, url_to_full_document, expected_source_id, None
            )

            # Verify the provided source_id was used
            assert len(source_ids_seen) == 1
            assert source_ids_seen[0] == expected_source_id
            # Verify it's NOT the domain
            assert "github.com" not in source_ids_seen[0]
            assert "python.org" not in source_ids_seen[0]
            assert "openai.com" not in source_ids_seen[0]

    def test_urlparse_not_imported(self):
        """Test that urlparse is not imported in code_extraction_service."""
        import src.server.services.crawling.code_extraction_service as module

        # Check that urlparse is not in the module's namespace
        assert not hasattr(module, "urlparse"), "urlparse should not be imported in code_extraction_service"

        # Check the module's actual imports
        import inspect

        source = inspect.getsource(module)
        assert "from urllib.parse import urlparse" not in source, (
            "Should not import urlparse since we don't extract domain from URL anymore"
        )