""" 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")