import pytest from cognee.modules.visualization.cognee_network_visualization import ( cognee_network_visualization, ) @pytest.mark.asyncio async def test_create_cognee_style_network_with_logo(tmp_path): nodes_data = [ ( 1, { "type": "Entity", "name": "Node1", "updated_at": 123, "created_at": 123, "source_task": "extract_graph_from_data", "source_pipeline": "cognify_pipeline", "source_node_set": "research_nodes", "source_user": "alice@example.com", }, ), ( 2, { "type": "DocumentChunk", "name": "Node2", "updated_at": 123, "created_at": 123, "source_task": "extract_chunks_from_documents", "source_pipeline": "cognify_pipeline", "source_node_set": "meeting_nodes", "source_user": "bob@example.com", }, ), ] edges_data = [ (1, 2, "related_to", {}), ] graph_data = (nodes_data, edges_data) html_output = await cognee_network_visualization( graph_data, str(tmp_path / "graph_visualization.html") ) assert isinstance(html_output, str) assert "users:id", "source_table": "posts", "target_table": "users", "relationship_type": "foreign_key", "source_column": "user_id", "target_column": "id", "description": "Posts user foreign key", }, ), ] edges_data = [ ("users", "database", "is_part_of", {"relationship_name": "is_part_of"}), ("posts", "database", "is_part_of", {"relationship_name": "is_part_of"}), ("posts", "posts_user_fk", "has_relationship", {"relationship_name": "foreign_key"}), ( "posts_user_fk", "users", "has_relationship", {"relationship_name": "foreign_key"}, ), ] html_output = await cognee_network_visualization( (nodes_data, edges_data), str(tmp_path / "schema_visualization.html") ) assert "const schemaGraphData = " in html_output assert '"type": "SchemaTable"' in html_output assert '"name": "users"' in html_output assert '"label": "foreign_key"' in html_output # Phase 2 redesign — old `graphDataToSchemaGraph` model builder # replaced by `buildSchemaModel`; renderer still exposed as # `window._renderSchemaGraph` so the tab-switch handler keeps working. assert "buildSchemaModel" in html_output assert "window._renderSchemaGraph" in html_output