library-desk/tests/test_entity_linking.py

"""
Comprehensive tests for Entity Linking system.

Tests cover:
- Finding entity mentions in pages
- Creating MENTIONS relationships in Neo4j
- Adding hyperlinks to wiki content
- Idempotency (safe to run multiple times)
- Protection of existing links (no nesting)
- Multi-tenancy isolation

Run with: pytest tests/test_entity_linking.py -v -s
"""

import pytest
import pytest_asyncio
from typing import AsyncGenerator

from src.clients.neo4j_client import Neo4jClient
from src.clients.wikijs_client import WikiJSClient
from src.services.graph_service import GraphService
from src.services.wiki_service import WikiService
from src.routers.entity_linking import (
    find_entity_mentions,
    add_entity_links_to_content,
    get_entities_with_paths
)
from src.config import get_settings

# Test user to isolate test data
TEST_USER = "entity-link-tester"


@pytest.fixture
def settings():
    """Get application settings."""
    return get_settings()


@pytest_asyncio.fixture
async def neo4j_client(settings, neo4j_test_uri) -> AsyncGenerator[Neo4jClient, None]:
    """Get connected Neo4j client."""
    client = Neo4jClient(
        uri=neo4j_test_uri,
        user=settings.neo4j_user,
        password=settings.neo4j_password
    )
    await client.connect()
    yield client
    await client.close()


@pytest_asyncio.fixture
async def wiki_client(wikijs_test_config) -> AsyncGenerator[WikiJSClient, None]:
    """Get Wiki.js client."""
    client = WikiJSClient(
        base_url=wikijs_test_config["base_url"],
        api_token=wikijs_test_config["api_token"]
    )
    yield client


@pytest_asyncio.fixture
async def graph_service(neo4j_client, wiki_client):
    """Get GraphService instance."""
    return GraphService(neo4j_client, wiki_client)


@pytest_asyncio.fixture
async def wiki_service(wiki_client):
    """Get WikiService instance."""
    return WikiService(wiki_client)


@pytest_asyncio.fixture
async def test_entities(graph_service):
    """Create test entities in graph."""
    from src.core.multi_tenancy import get_neo4j_user_base_label

    user_base_label = get_neo4j_user_base_label(TEST_USER)

    # Clean up any existing test entities
    cleanup_query = f"""
    MATCH (n:{user_base_label})
    WHERE n.name IN ['Docker', 'Kubernetes', 'PostgreSQL']
    DETACH DELETE n
    """
    await graph_service.neo4j.execute_query(cleanup_query)

    # Create test entities
    create_query = f"""
    CREATE (d:{user_base_label}:Technology {{name: 'Docker', type: 'technology'}})
    CREATE (k:{user_base_label}:Technology {{name: 'Kubernetes', type: 'technology'}})
    CREATE (p:{user_base_label}:Technology {{name: 'PostgreSQL', type: 'technology'}})
    RETURN d.name, k.name, p.name
    """
    await graph_service.neo4j.execute_query(create_query)

    yield ["Docker", "Kubernetes", "PostgreSQL"]

    # Cleanup after test
    await graph_service.neo4j.execute_query(cleanup_query)


# ============================================================================
# Unit Tests - Entity Mention Detection
# ============================================================================

class TestFindEntityMentions:
    """Test finding entity mentions in content."""

    def test_find_single_mention(self):
        """Test finding a single entity mention."""
        content = "Docker is a containerization platform."
        entities = [
            {"name": "Docker", "type": "technology"}
        ]

        found = find_entity_mentions(content, entities)

        assert len(found) == 1
        assert found[0]["name"] == "Docker"
        assert found[0]["mentions"] == 1

    def test_find_multiple_mentions(self):
        """Test finding multiple mentions of same entity."""
        content = "Docker containers run on Docker Engine. Docker is great!"
        entities = [
            {"name": "Docker", "type": "technology"}
        ]

        found = find_entity_mentions(content, entities)

        assert len(found) == 1
        assert found[0]["name"] == "Docker"
        assert found[0]["mentions"] == 3

    def test_case_insensitive_matching(self):
        """Test case-insensitive entity matching."""
        content = "docker and DOCKER and Docker are the same"
        entities = [
            {"name": "Docker", "type": "technology"}
        ]

        found = find_entity_mentions(content, entities)

        assert len(found) == 1
        assert found[0]["mentions"] == 3

    def test_whole_word_matching(self):
        """Test that partial word matches are excluded."""
        content = "Kubernetes and Kubernetes-based and MyKubernetes"
        entities = [
            {"name": "Kubernetes", "type": "technology"}
        ]

        found = find_entity_mentions(content, entities)

        assert len(found) == 1
        # Regex \b matches at hyphens, so "Kubernetes-based" contains "Kubernetes"
        # Only "MyKubernetes" is excluded (no word boundary)
        assert found[0]["mentions"] == 2  # "Kubernetes" and "Kubernetes-based"

    def test_ignore_short_names(self):
        """Test that entities with names <3 chars are ignored."""
        content = "Go is a programming language by Google"
        entities = [
            {"name": "Go", "type": "language"},  # Too short
            {"name": "Google", "type": "organization"}
        ]

        found = find_entity_mentions(content, entities)

        assert len(found) == 1
        assert found[0]["name"] == "Google"

    def test_sort_by_mention_count(self):
        """Test results are sorted by mention count."""
        content = "Docker Docker Docker. Kubernetes Kubernetes. PostgreSQL."
        entities = [
            {"name": "PostgreSQL", "type": "database"},
            {"name": "Docker", "type": "technology"},
            {"name": "Kubernetes", "type": "technology"}
        ]

        found = find_entity_mentions(content, entities)

        assert len(found) == 3
        assert found[0]["name"] == "Docker"  # Most mentions
        assert found[0]["mentions"] == 3
        assert found[1]["name"] == "Kubernetes"
        assert found[1]["mentions"] == 2
        assert found[2]["name"] == "PostgreSQL"
        assert found[2]["mentions"] == 1


# ============================================================================
# Unit Tests - Content Link Addition
# ============================================================================

class TestAddEntityLinksToContent:
    """Test adding hyperlinks to content."""

    def test_add_single_link(self):
        """Test adding a single entity link."""
        content = "Docker is a containerization platform."
        entities = [
            {"name": "Docker", "path": "users/test/docker"}
        ]

        updated, count = add_entity_links_to_content(content, entities)

        assert count == 1
        assert "[Docker](/users/test/docker)" in updated

    def test_add_multiple_instances(self):
        """Test linking all instances of an entity."""
        content = "Docker containers run on Docker Engine."
        entities = [
            {"name": "Docker", "path": "users/test/docker"}
        ]

        updated, count = add_entity_links_to_content(content, entities)

        assert count == 2  # Both instances linked
        assert updated.count("[Docker](/users/test/docker)") == 2

    def test_skip_entities_without_path(self):
        """Test that entities without wiki pages are not linked."""
        content = "Docker and Kubernetes are used together."
        entities = [
            {"name": "Docker", "path": "users/test/docker"},
            {"name": "Kubernetes", "path": None}  # No page
        ]

        updated, count = add_entity_links_to_content(content, entities)

        assert count == 1  # Only Docker
        assert "[Docker](/users/test/docker)" in updated
        assert "[Kubernetes]" not in updated

    def test_protect_existing_links(self):
        """Test that existing markdown links are not modified."""
        content = "See [Docker](https://docker.com) for more info. Docker is great!"
        entities = [
            {"name": "Docker", "path": "users/test/docker"}
        ]

        updated, count = add_entity_links_to_content(content, entities)

        # Should link the second "Docker" but not the one already linked
        assert count == 1
        assert "[Docker](https://docker.com)" in updated  # Preserved
        assert updated.count("[Docker](/users/test/docker)") == 1

    def test_no_nested_links(self):
        """Test that entity names in URLs are not linked."""
        content = "Check [Docker Hub](/docker/hub) for images."
        entities = [
            {"name": "Docker", "path": "users/test/docker"}
        ]

        updated, count = add_entity_links_to_content(content, entities)

        # "Docker" in the URL path should not be linked
        assert count == 0
        assert "[Docker Hub](/docker/hub)" in updated  # Unchanged

    def test_longest_first_matching(self):
        """Test that longer entity names are matched first."""
        content = "Machine Learning and Machine are different."
        entities = [
            {"name": "Machine Learning", "path": "users/test/ml"},
            {"name": "Machine", "path": "users/test/machine"}
        ]

        updated, count = add_entity_links_to_content(content, entities)

        # Should link "Machine Learning" first, leaving "Machine" alone
        assert "[Machine Learning](/users/test/ml)" in updated
        assert count >= 1


# ============================================================================
# Integration Tests - Full Entity Linking Flow
# ============================================================================

class TestEntityLinkingIntegration:
    """Test full entity linking flow."""

    @pytest.mark.asyncio
    async def test_get_entities_with_paths(self, graph_service, test_entities):
        """Test retrieving entities and their wiki page paths."""
        entities = await get_entities_with_paths(graph_service, TEST_USER)

        # Should find our test entities
        entity_names = [e["name"] for e in entities]
        assert "Docker" in entity_names
        assert "Kubernetes" in entity_names
        assert "PostgreSQL" in entity_names

    @pytest.mark.asyncio
    async def test_create_mentions_relationships(self, graph_service, test_entities):
        """Test creating MENTIONS relationships."""
        from src.core.multi_tenancy import get_neo4j_user_label

        user_doc_label = get_neo4j_user_label(TEST_USER)

        # Create a test document node
        doc_query = f"""
        CREATE (d:{user_doc_label}:Document {{
            page_id: 9999,
            title: 'Test Doc',
            path: 'users/test/doc'
        }})
        RETURN d
        """
        await graph_service.neo4j.execute_query(doc_query)

        # Create MENTIONS relationships
        found_entities = [
            {"name": "Docker"},
            {"name": "Kubernetes"}
        ]

        new_links = await graph_service.create_entity_mentions(
            page_id=9999,
            user=TEST_USER,
            entity_names=found_entities
        )

        assert new_links == 2

        # Verify relationships exist
        verify_query = f"""
        MATCH (d:{user_doc_label}:Document {{page_id: 9999}})-[r:MENTIONS]->(e)
        RETURN count(r) as mention_count
        """
        result = await graph_service.neo4j.execute_query(verify_query)
        assert result[0]["mention_count"] == 2

        # Cleanup
        cleanup_query = f"""
        MATCH (d:{user_doc_label}:Document {{page_id: 9999}})
        DETACH DELETE d
        """
        await graph_service.neo4j.execute_query(cleanup_query)

    @pytest.mark.asyncio
    async def test_idempotency(self, graph_service):
        """Test that entity linking is idempotent."""
        from src.core.multi_tenancy import get_neo4j_user_label, get_neo4j_user_base_label

        user_doc_label = get_neo4j_user_label(TEST_USER)
        user_base_label = get_neo4j_user_base_label(TEST_USER)

        # Aggressively clean up ALL test data first (fresh start)
        cleanup_all = f"""
        MATCH (n)
        WHERE (n:{user_base_label} OR n:{user_doc_label})
          AND (n.page_id = 9998 OR n.name = 'TestDockerEntity')
        DETACH DELETE n
        """
        await graph_service.neo4j.execute_query(cleanup_all)

        # Create a unique test entity
        entity_query = f"""
        CREATE (e:{user_base_label}:Technology {{name: 'TestDockerEntity', type: 'technology'}})
        RETURN e
        """
        await graph_service.neo4j.execute_query(entity_query)

        # Create test document
        doc_query = f"""
        CREATE (d:{user_doc_label}:Document {{
            page_id: 9998,
            title: 'Test Doc 2',
            path: 'users/test/doc2'
        }})
        RETURN d
        """
        await graph_service.neo4j.execute_query(doc_query)

        found_entities = [{"name": "TestDockerEntity"}]

        # Link once
        first_run = await graph_service.create_entity_mentions(
            page_id=9998,
            user=TEST_USER,
            entity_names=found_entities
        )
        assert first_run == 1

        # Link again - should not create duplicates
        second_run = await graph_service.create_entity_mentions(
            page_id=9998,
            user=TEST_USER,
            entity_names=found_entities
        )
        assert second_run == 0  # No new links

        # Verify only one relationship exists
        verify_query = f"""
        MATCH (d:{user_doc_label}:Document {{page_id: 9998}})-[r:MENTIONS]->()
        RETURN count(r) as mention_count
        """
        result = await graph_service.neo4j.execute_query(verify_query)
        assert result[0]["mention_count"] == 1

        # Cleanup
        cleanup_query = f"""
        MATCH (n)
        WHERE (n:{user_base_label} OR n:{user_doc_label})
          AND (n.page_id = 9998 OR n.name = 'TestDockerEntity')
        DETACH DELETE n
        """
        await graph_service.neo4j.execute_query(cleanup_query)


# ============================================================================
# Multi-Tenancy Tests
# ============================================================================

class TestEntityLinkingMultiTenancy:
    """Test multi-tenancy isolation in entity linking."""

    @pytest.mark.asyncio
    async def test_user_isolation(self, graph_service):
        """Test that entities are isolated by user."""
        from src.core.multi_tenancy import get_neo4j_user_base_label

        user1_label = get_neo4j_user_base_label("user1")
        user2_label = get_neo4j_user_base_label("user2")

        # Create entity for user1
        create_user1 = f"""
        CREATE (e:{user1_label}:Technology {{name: 'Docker', type: 'technology'}})
        RETURN e
        """
        await graph_service.neo4j.execute_query(create_user1)

        # Create entity for user2
        create_user2 = f"""
        CREATE (e:{user2_label}:Technology {{name: 'Docker', type: 'technology'}})
        RETURN e
        """
        await graph_service.neo4j.execute_query(create_user2)

        # Get entities for user1 - should only see user1's entities
        entities_user1 = await get_entities_with_paths(graph_service, "user1")
        entity_names_user1 = [e["name"] for e in entities_user1]

        # Verify isolation
        assert "Docker" in entity_names_user1
        # We can't verify the exact count without knowing what else is in the DB,
        # but we verified we can retrieve entities for user1

        # Cleanup
        await graph_service.neo4j.execute_query(f"MATCH (e:{user1_label}) WHERE e.name = 'Docker' DETACH DELETE e")
        await graph_service.neo4j.execute_query(f"MATCH (e:{user2_label}) WHERE e.name = 'Docker' DETACH DELETE e")


# ============================================================================
# Cleanup
# ============================================================================

@pytest.mark.asyncio
async def test_cleanup_entity_linking_test_data(neo4j_client):
    """Clean up all test data created by entity linking tests."""
    from src.core.multi_tenancy import get_neo4j_user_base_label

    for user in [TEST_USER, "user1", "user2"]:
        user_label = get_neo4j_user_base_label(user)

        cleanup_query = f"""
        MATCH (n:{user_label})
        WHERE n.page_id IN [9999, 9998]
           OR n.name IN ['Docker', 'Kubernetes', 'PostgreSQL']
        DETACH DELETE n
        """
        await neo4j_client.execute_query(cleanup_query)

    print("\n✓ Cleaned up entity linking test data")