diff --git a/.gitignore b/.gitignore index 89ef3acf..f4a119d5 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,5 @@ venv/ .ruff_cache/ .vscode/ .coverage +config.json __pycache__/ \ No newline at end of file diff --git a/prometheus/configuration/__init__.py b/prometheus/configuration/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/prometheus/configuration/config.py b/prometheus/configuration/config.py new file mode 100644 index 00000000..6316c3d6 --- /dev/null +++ b/prometheus/configuration/config.py @@ -0,0 +1,12 @@ +import json +from pathlib import Path + +CONFIG_FILE = Path(__file__).parent / "config.json" + + +def load_config(): + with open(CONFIG_FILE, "r") as f: + return json.load(f) + + +config = load_config() diff --git a/prometheus/configuration/example_config.json b/prometheus/configuration/example_config.json new file mode 100644 index 00000000..9497f829 --- /dev/null +++ b/prometheus/configuration/example_config.json @@ -0,0 +1,12 @@ +{ + "neo4j": { + "uri": "CHANGE_ME", + "username": "CHANGE_ME", + "password": "CHANGE_ME", + "database": "neo4j", + "batch_size": 1000 + }, + "knowledge_graph": { + "max_ast_depth": 5 + } +} \ No newline at end of file diff --git a/prometheus/graph/graph_types.py b/prometheus/graph/graph_types.py index b60bed80..94dd0389 100644 --- a/prometheus/graph/graph_types.py +++ b/prometheus/graph/graph_types.py @@ -24,8 +24,8 @@ class ASTNode: Attributes: type: The tree-sitter node type. - start_line: The starting line number. - end_line: The ending line number. + start_line: The starting line number. 0-indexed and inclusive. + end_line: The ending line number. 0-indexed and inclusive. text: The source code correcpsonding to the node. """ @@ -79,7 +79,9 @@ def to_neo4j_node(self) -> Union["Neo4jFileNode", "Neo4jASTNode", "Neo4jTextNode ) case TextNode(): return Neo4jTextNode( - node_id=self.node_id, text=self.node.text, metadata=self.node.metadata + node_id=self.node_id, + text=self.node.text, + metadata=self.node.metadata, ) case _: raise ValueError("Unknown KnowledgeGraphNode.node type") @@ -122,23 +124,28 @@ def to_neo4j_edge( match self.type: case KnowledgeGraphEdgeType.has_file: return Neo4jHasFileEdge( - source=self.source.to_neo4j_node(), target=self.target.to_neo4j_node() + source=self.source.to_neo4j_node(), + target=self.target.to_neo4j_node(), ) case KnowledgeGraphEdgeType.has_ast: return Neo4jHasASTEdge( - source=self.source.to_neo4j_node(), target=self.target.to_neo4j_node() + source=self.source.to_neo4j_node(), + target=self.target.to_neo4j_node(), ) case KnowledgeGraphEdgeType.parent_of: return Neo4jParentOfEdge( - source=self.source.to_neo4j_node(), target=self.target.to_neo4j_node() + source=self.source.to_neo4j_node(), + target=self.target.to_neo4j_node(), ) case KnowledgeGraphEdgeType.has_text: return Neo4jHasTextEdge( - source=self.source.to_neo4j_node(), target=self.target.to_neo4j_node() + source=self.source.to_neo4j_node(), + target=self.target.to_neo4j_node(), ) case KnowledgeGraphEdgeType.next_chunk: return Neo4jNextChunkEdge( - source=self.source.to_neo4j_node(), target=self.target.to_neo4j_node() + source=self.source.to_neo4j_node(), + target=self.target.to_neo4j_node(), ) case _: raise ValueError(f"Unknown edge type: {self.type}") diff --git a/prometheus/graph/knowledge_graph.py b/prometheus/graph/knowledge_graph.py index aca78cf1..41fbebbd 100644 --- a/prometheus/graph/knowledge_graph.py +++ b/prometheus/graph/knowledge_graph.py @@ -84,14 +84,16 @@ def _build_graph(self, root_dir: Path): for child_file in sorted(file.iterdir()): child_file_node = FileNode( basename=child_file.name, - relative_path=str(child_file.relative_to(root_dir)), + relative_path=str(child_file.relative_to(root_dir).as_posix()), ) kg_child_file_node = KnowledgeGraphNode(self._next_node_id, child_file_node) self._next_node_id += 1 self._knowledge_graph_nodes.append(kg_child_file_node) self._knowledge_graph_edges.append( KnowledgeGraphEdge( - kg_file_path_node, kg_child_file_node, KnowledgeGraphEdgeType.has_file + kg_file_path_node, + kg_child_file_node, + KnowledgeGraphEdgeType.has_file, ) ) diff --git a/prometheus/tools/__init__.py b/prometheus/tools/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/prometheus/tools/graph_traversal.py b/prometheus/tools/graph_traversal.py new file mode 100644 index 00000000..1ba7c054 --- /dev/null +++ b/prometheus/tools/graph_traversal.py @@ -0,0 +1,240 @@ +from pathlib import Path +from pydantic import BaseModel, Field + +from neo4j import GraphDatabase + +from prometheus.parser import tree_sitter_parser +from prometheus.utils import neo4j_util + +MAX_RESULT = 20 + +############################################################################### +# FileNode retrieval # +############################################################################### + + +class FindFileNodeWithBasenameInput(BaseModel): + basename: str = Field("The basename of FileNode to search for") + + +def find_file_node_with_basename(basename: str, driver: GraphDatabase.driver) -> str: + query = f""" + MATCH (f:FileNode {{ basename: '{basename}' }}) + RETURN f AS FileNode + ORDER BY f.node_id + LIMIT {MAX_RESULT} + """ + return neo4j_util.run_neo4j_query(query, driver) + + +class FindFileNodeWithRelativePathInput(BaseModel): + relative_path: str = Field("The relative_path of FileNode to search for") + + +def find_file_node_with_relative_path( + relative_path: str, driver: GraphDatabase.driver +) -> str: + query = f"""\ + MATCH (f:FileNode {{ relative_path: '{relative_path}' }}) + RETURN f AS FileNode + ORDER BY f.node_id + LIMIT {MAX_RESULT} + """ + return neo4j_util.run_neo4j_query(query, driver) + + +############################################################################### +# ASTNode retrieval # +############################################################################### + + +class FindASTNodeWithTextInput(BaseModel): + text: str = Field("Search ASTNode that exactly contains this text.") + + +def find_ast_node_with_text(text: str, driver: GraphDatabase.driver) -> str: + query = f"""\ + MATCH (f:FileNode) -[:HAS_AST]-> (:ASTNode) -[:PARENT_OF*]-> (a:ASTNode) + WHERE a.text CONTAINS '{text}' + RETURN f as FileNode, a AS ASTNode + ORDER BY SIZE(a.text) + LIMIT {MAX_RESULT} + """ + return neo4j_util.run_neo4j_query(query, driver) + + +class FindASTNodeWithTypeInput(BaseModel): + type: str = Field("Search ASTNode that has this tree-sitter node type.") + + +def find_ast_node_with_type(type: str, driver: GraphDatabase.driver) -> str: + query = f"""\ + MATCH (f:FileNode) -[:HAS_AST]-> (:ASTNode) -[:PARENT_OF*]-> (a:ASTNode {{ type: '{type}' }}) + RETURN f as FileNode, a AS ASTNode + ORDER BY a.node_id + LIMIT {MAX_RESULT} + """ + return neo4j_util.run_neo4j_query(query, driver) + + +class FindASTNodeWithTextInFileInput(BaseModel): + text: str = Field("Search ASTNode that exactly contains this text.") + basename: str = Field("The basename of FileNode to search ASTNode.") + + +def find_ast_node_with_text_in_file( + text: str, basename: str, driver: GraphDatabase.driver +) -> str: + query = f"""\ + MATCH (f:FileNode) -[:HAS_AST]-> (:ASTNode) -[:PARENT_OF*]-> (a:ASTNode) + WHERE f.basename = '{basename}' AND a.text CONTAINS '{text}' + RETURN f as FileNode, a AS ASTNode + ORDER BY SIZE(a.text) + LIMIT {MAX_RESULT} + """ + return neo4j_util.run_neo4j_query(query, driver) + + +class FindASTNodeWithTypeInFileInput(BaseModel): + type: str = Field("Search ASTNode with this tree-sitter node type.") + basename: str = Field("The basename of FileNode to search ASTNode.") + + +def find_ast_node_with_type_in_file( + type: str, basename: str, driver: GraphDatabase.driver +) -> str: + query = f"""\ + MATCH (f:FileNode) -[:HAS_AST]-> (:ASTNode) -[:PARENT_OF*]-> (a:ASTNode) + WHERE f.basename = '{basename}' AND a.type = '{type}' + RETURN f as FileNode, a AS ASTNode + ORDER BY SIZE(a.text) + LIMIT {MAX_RESULT} + """ + return neo4j_util.run_neo4j_query(query, driver) + + +class FindASTNodeWithTypeAndTextInput(BaseModel): + type: str = Field("Search ASTNode with this tree-sitter node type.") + text: str = Field("Search ASTNode that exactly contains this text.") + + +def find_ast_node_with_type_and_text( + type: str, text: str, driver: GraphDatabase.driver +) -> str: + query = f"""\ + MATCH (f:FileNode) -[:HAS_AST]-> (:ASTNode) -[:PARENT_OF*]-> (a:ASTNode) + WHERE a.type = '{type}' AND a.text CONTAINS '{text}' + RETURN f as FileNode, a AS ASTNode + ORDER BY SIZE(a.text) + LIMIT {MAX_RESULT} + """ + return neo4j_util.run_neo4j_query(query, driver) + + +############################################################################### +# TextNode retrieval # +############################################################################### + + +class FindTextNodeWithTextInput(BaseModel): + text: str = Field("Search TextNode that exactly contains this text.") + + +def find_text_node_with_text(text: str, driver: GraphDatabase.driver) -> str: + query = f"""\ + MATCH (f:FileNode) -[:HAS_TEXT]-> (t:TextNode) + WHERE t.text CONTAINS '{text}' + RETURN f as FileNode, t AS TextNode + ORDER BY t.node_id + LIMIT {MAX_RESULT} + """ + return neo4j_util.run_neo4j_query(query, driver) + + +class FindTextNodeWithTextInFileInput(BaseModel): + text: str = Field("Search TextNode that exactly contains this text.") + basename: str = Field("The basename of FileNode to search TextNode.") + + +def find_text_node_with_text_in_file( + text: str, basename: str, driver: GraphDatabase.driver +) -> str: + query = f"""\ + MATCH (f:FileNode) -[:HAS_TEXT]-> (t:TextNode) + WHERE f.basename = '{basename}' AND t.text CONTAINS '{text}' + RETURN f as FileNode, t AS TextNode + ORDER BY t.node_id + LIMIT {MAX_RESULT} + """ + return neo4j_util.run_neo4j_query(query, driver) + + +class GetNextTextNodeWithNodeIdInput(BaseModel): + node_id: int = Field("Get the next TextNode of this given node_id.") + + +def get_next_text_node_with_node_id(node_id: str, driver: GraphDatabase.driver) -> str: + query = f"""\ + MATCH (a:TextNode {{ node_id: {node_id} }}) -[:NEXT_CHUNK]-> (b:TextNode) + RETURN b as TextNode + """ + return neo4j_util.run_neo4j_query(query, driver) + + +############################################################################### +# Other # +############################################################################### + + +class PreviewFileContentWithBasenameInput(BaseModel): + basename: str = Field("The basename of FileNode to preview.") + + +def preview_file_content_with_basename( + basename: str, driver: GraphDatabase.driver +) -> str: + source_code_query = f"""\ + MATCH (f:FileNode {{ basename: '{basename}' }}) -[:HAS_AST]-> (a:ASTNode) + WITH f, apoc.text.split(a.text, '\\R') AS lines + RETURN f as FileNode, apoc.text.join(lines[0..300], '\\n') AS preview + ORDER BY f.node_id + """ + + text_query = f"""\ + MATCH (f:FileNode {{ basename: '{basename}' }}) -[:HAS_TEXT]-> (t:TextNode) + WHERE NOT EXISTS((:TextNode) -[:NEXT_CHUNK]-> (t)) + RETURN f as FileNode, t.text AS preview + ORDER BY f.node_id + """ + + if tree_sitter_parser.supports_file(Path(basename)): + return neo4j_util.run_neo4j_query(source_code_query, driver) + return neo4j_util.run_neo4j_query(text_query, driver) + + +class GetParentNodeInput(BaseModel): + node_id: str = Field(description="Get parent node of node with this node_id") + + +def get_parent_node(node_id: int, driver: GraphDatabase.driver) -> str: + query = f"""\ + MATCH (p) -[r]-> (c {{ node_id: {node_id} }}) + WHERE type(r) IN ['HAS_FILE', 'HAS_TEXT', 'HAS_AST', 'PARENT_OF'] + RETURN p as ParentNode, head(labels(p)) as ParentNodeType + ORDER BY p.node_id + """ + return neo4j_util.run_neo4j_query(query, driver) + + +class GetChildrenNodeInput(BaseModel): + node_id: str = Field(description="Get children nodes of node with this node_id") + + +def get_children_node(node_id: int, driver: GraphDatabase.driver) -> str: + query = f"""\ + MATCH (p {{ node_id: {node_id} }}) -[r]-> (c) + WHERE type(r) IN ['HAS_FILE', 'HAS_TEXT', 'HAS_AST', 'PARENT_OF'] + RETURN c as ChildNode, head(labels(p)) as ChildNodeType + ORDER BY c.node_id + """ + return neo4j_util.run_neo4j_query(query, driver) diff --git a/prometheus/utils/__init__.py b/prometheus/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/prometheus/utils/neo4j_util.py b/prometheus/utils/neo4j_util.py new file mode 100644 index 00000000..f4ff144d --- /dev/null +++ b/prometheus/utils/neo4j_util.py @@ -0,0 +1,39 @@ +import neo4j + + +def format_neo4j_result(result: neo4j.Result) -> str: + """Format a Neo4j result into a string. + + Args: + result: The result from a Neo4j query. + + Returns: + A string representation of the result. + """ + data = result.data() + output = "" + for index, row_result in enumerate(data): + output += f"Result {index+1}:\n" + for key in sorted(row_result.keys()): + output += f"{key}: {row_result[key]}\n" + output += "\n\n" + return output.strip() + + +def run_neo4j_query(query: str, driver: neo4j.GraphDatabase.driver) -> str: + """Run a read-only Neo4j query and format the result into a string. + + Args: + query: The query to run. + driver: The Neo4j driver to use. + + Returns: + A string representation of the result. + """ + + def query_transaction(tx): + result = tx.run(query) + return format_neo4j_result(result) + + with driver.session() as session: + return session.execute_read(query_transaction) diff --git a/tests/graph/test_knowledge_graph.py b/tests/graph/test_knowledge_graph.py index 86db20d0..e46419b4 100644 --- a/tests/graph/test_knowledge_graph.py +++ b/tests/graph/test_knowledge_graph.py @@ -5,17 +5,17 @@ def test_build_graph(): knowledge_graph = KnowledgeGraph(test_project_paths.TEST_PROJECT_PATH, 1000) - assert knowledge_graph._next_node_id == 97 + assert knowledge_graph._next_node_id == 96 # 8 FileNode - # 85 ASTnode + # 84 ASTnode # 4 TextNode - assert len(knowledge_graph._knowledge_graph_nodes) == 97 - assert len(knowledge_graph._knowledge_graph_edges) == 99 + assert len(knowledge_graph._knowledge_graph_nodes) == 96 + assert len(knowledge_graph._knowledge_graph_edges) == 98 assert len(knowledge_graph.get_file_nodes()) == 8 - assert len(knowledge_graph.get_ast_nodes()) == 85 + assert len(knowledge_graph.get_ast_nodes()) == 84 assert len(knowledge_graph.get_text_nodes()) == 4 - assert len(knowledge_graph.get_parent_of_edges()) == 82 + assert len(knowledge_graph.get_parent_of_edges()) == 81 assert len(knowledge_graph.get_has_file_edges()) == 7 assert len(knowledge_graph.get_has_ast_edges()) == 3 assert len(knowledge_graph.get_has_text_edges()) == 4 diff --git a/tests/neo4j/test_handler.py b/tests/neo4j/test_handler.py index 8a0b3bea..615a802f 100644 --- a/tests/neo4j/test_handler.py +++ b/tests/neo4j/test_handler.py @@ -15,9 +15,10 @@ @pytest.fixture(scope="session") def setup_neo4j_container(): kg = KnowledgeGraph(test_project_paths.TEST_PROJECT_PATH, 1000) - with Neo4jContainer( + container = Neo4jContainer( image=NEO4J_IMAGE, username=NEO4J_USERNAME, password=NEO4J_PASSWORD - ) as neo4j_container: + ).with_env("NEO4JLABS_PLUGINS", '["apoc"]') + with container as neo4j_container: uri = neo4j_container.get_connection_url() handler = Handler(uri, NEO4J_USERNAME, NEO4J_PASSWORD, "neo4j", 100) handler.write_knowledge_graph(kg) @@ -40,7 +41,7 @@ def _count_num_ast_nodes(tx): with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: with driver.session() as session: read_ast_nodes = session.execute_read(_count_num_ast_nodes) - assert len(read_ast_nodes) == 85 + assert len(read_ast_nodes) == 84 def test_num_file_nodes(setup_neo4j_container): @@ -94,7 +95,7 @@ def _count_num_parent_of_edges(tx): with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: with driver.session() as session: read_parent_of_edges = session.execute_read(_count_num_parent_of_edges) - assert len(read_parent_of_edges) == 82 + assert len(read_parent_of_edges) == 81 def test_num_has_file_edges(setup_neo4j_container): diff --git a/tests/test_project/test.c b/tests/test_project/test.c index 822faf74..79a11609 100644 --- a/tests/test_project/test.c +++ b/tests/test_project/test.c @@ -1,6 +1,6 @@ #include int main() { - printf("Hello world!\n"); + printf("Hello world!"); return 0; } \ No newline at end of file diff --git a/tests/tools/__init__.py b/tests/tools/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/tools/test_graph_traversal.py b/tests/tools/test_graph_traversal.py new file mode 100644 index 00000000..027e6faf --- /dev/null +++ b/tests/tools/test_graph_traversal.py @@ -0,0 +1,322 @@ +from neo4j import GraphDatabase +import pytest + +from testcontainers.neo4j import Neo4jContainer + +from prometheus.graph.knowledge_graph import KnowledgeGraph +from prometheus.neo4j.handler import Handler +from prometheus.tools import graph_traversal +from tests.test_utils import test_project_paths + +NEO4J_IMAGE = "neo4j:5.20.0" +NEO4J_USERNAME = "neo4j" +NEO4J_PASSWORD = "password" + + +@pytest.fixture(scope="session") +def setup_neo4j_container(): + kg = KnowledgeGraph(test_project_paths.TEST_PROJECT_PATH, 1000) + container = Neo4jContainer( + image=NEO4J_IMAGE, username=NEO4J_USERNAME, password=NEO4J_PASSWORD + ).with_env("NEO4JLABS_PLUGINS", '["apoc"]') + with container as neo4j_container: + uri = neo4j_container.get_connection_url() + handler = Handler(uri, NEO4J_USERNAME, NEO4J_PASSWORD, "neo4j", 100) + handler.write_knowledge_graph(kg) + handler.close() + yield neo4j_container + + +def test_find_file_node_with_basename(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.find_file_node_with_basename( + test_project_paths.PYTHON_FILE.name, driver + ) + + basename = test_project_paths.PYTHON_FILE.name + relative_path = str( + test_project_paths.PYTHON_FILE.relative_to( + test_project_paths.TEST_PROJECT_PATH + ).as_posix() + ) + assert result.count("FileNode") == 1 + assert f"'basename': '{basename}'" in result + assert f"'relative_path': '{relative_path}'" in result + + +def test_find_file_node_with_relative_path(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + relative_path = str( + test_project_paths.MD_FILE.relative_to( + test_project_paths.TEST_PROJECT_PATH + ).as_posix() + ) + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.find_file_node_with_relative_path(relative_path, driver) + + basename = test_project_paths.MD_FILE.name + assert result.count("FileNode") == 1 + assert f"'basename': '{basename}'" in result + assert f"'relative_path': '{relative_path}'" in result + + +def test_find_ast_node_with_text(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + text = "System.out.println" + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.find_ast_node_with_text(text, driver) + + assert "FileNode" in result + assert "ASTNode" in result + + basename = test_project_paths.JAVA_FILE.name + relative_path = str( + test_project_paths.JAVA_FILE.relative_to( + test_project_paths.TEST_PROJECT_PATH + ).as_posix() + ) + assert f"'basename': '{basename}'" in result + assert f"'relative_path': '{relative_path}'" in result + assert "'text': 'System.out.println(\"Hello world!\")'" in result + assert "'type': 'method_invocation'" in result + assert "'start_line': 2" in result + assert "'end_line': 2" in result + + +def test_find_ast_node_with_type(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + type = "argument_list" + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.find_ast_node_with_type(type, driver) + + assert "FileNode" in result + assert "ASTNode" in result + + basename = test_project_paths.JAVA_FILE.name + relative_path = str( + test_project_paths.JAVA_FILE.relative_to( + test_project_paths.TEST_PROJECT_PATH + ).as_posix() + ) + assert f"'basename': '{basename}'" in result + assert f"'relative_path': '{relative_path}'" in result + assert "'text': '(\"Hello world!\")'" in result + assert f"'type': '{type}'" in result + assert "'start_line': 2" in result + assert "'end_line': 2" in result + + +def test_find_ast_node_with_text_in_file(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + text = "printf" + basename = test_project_paths.C_FILE.name + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.find_ast_node_with_text_in_file(text, basename, driver) + + relative_path = str( + test_project_paths.C_FILE.relative_to( + test_project_paths.TEST_PROJECT_PATH + ).as_posix() + ) + assert f"'basename': '{basename}'" in result + assert f"'relative_path': '{relative_path}'" in result + assert f"'text': '{text}'" in result + assert "'type': 'identifier'" in result + assert "'start_line': 3" in result + assert "'end_line': 3" in result + + +def test_find_ast_node_with_type_in_file(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + type = "string_literal" + basename = test_project_paths.C_FILE.name + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.find_ast_node_with_type_in_file(type, basename, driver) + + relative_path = str( + test_project_paths.C_FILE.relative_to( + test_project_paths.TEST_PROJECT_PATH + ).as_posix() + ) + assert f"'basename': '{basename}'" in result + assert f"'relative_path': '{relative_path}'" in result + assert "'text': '\"Hello world!\"'" in result + assert f"'type': '{type}'" in result + assert "'start_line': 3" in result + assert "'end_line': 3" in result + + +def test_find_ast_node_with_type_and_text(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + type = "string_literal" + text = "Hello world!" + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.find_ast_node_with_type_and_text(type, text, driver) + + basename = test_project_paths.C_FILE.name + relative_path = str( + test_project_paths.C_FILE.relative_to( + test_project_paths.TEST_PROJECT_PATH + ).as_posix() + ) + assert f"'basename': '{basename}'" in result + assert f"'relative_path': '{relative_path}'" in result + assert f"'text': '\"{text}\"'" in result + assert f"'type': '{type}'" in result + assert "'start_line': 3" in result + assert "'end_line': 3" in result + + +def test_find_text_node_with_text(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + text = "Text under header A." + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.find_text_node_with_text(text, driver) + + assert "FileNode" in result + assert "TextNode" in result + + basename = test_project_paths.MD_FILE.name + relative_path = str( + test_project_paths.MD_FILE.relative_to( + test_project_paths.TEST_PROJECT_PATH + ).as_posix() + ) + assert f"'basename': '{basename}'" in result + assert f"'relative_path': '{relative_path}'" in result + assert f"'text': '{text}'" in result + assert "'metadata': \"{'Header 1': 'A'}\"" in result + + +def test_find_text_node_with_text_in_file(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + text = "Text under header B." + basename = test_project_paths.MD_FILE.name + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.find_text_node_with_text_in_file(text, basename, driver) + + assert "FileNode" in result + assert "TextNode" in result + + relative_path = str( + test_project_paths.MD_FILE.relative_to( + test_project_paths.TEST_PROJECT_PATH + ).as_posix() + ) + assert f"'basename': '{basename}'" in result + assert f"'relative_path': '{relative_path}'" in result + assert f"'text': '{text}'" in result + assert "'metadata': \"{'Header 1': 'A', 'Header 2': 'B'}\"" in result + + +def test_get_next_text_node_with_node_id(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + # node_id of TextNode 'Text under header B.' + node_id = 36 + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.get_next_text_node_with_node_id(node_id, driver) + + assert "'text': 'Text under header C.'" in result + assert "'metadata': \"{'Header 1': 'A', 'Header 2': 'C'}\"" in result + + +def test_preview_source_code_file_content_with_basename(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + basename = test_project_paths.C_FILE.name + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.preview_file_content_with_basename(basename, driver) + + assert "FileNode" in result + assert "preview" in result + + source_code = test_project_paths.C_FILE.open().read() + basename = test_project_paths.C_FILE.name + relative_path = str( + test_project_paths.C_FILE.relative_to( + test_project_paths.TEST_PROJECT_PATH + ).as_posix() + ) + assert f"'basename': '{basename}'" in result + assert f"'relative_path': '{relative_path}'" in result + assert source_code in result + + +def test_preview_text_file_content_with_basename(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + basename = test_project_paths.MD_FILE.name + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.preview_file_content_with_basename(basename, driver) + + assert "FileNode" in result + assert "preview" in result + + basename = test_project_paths.MD_FILE.name + relative_path = str( + test_project_paths.MD_FILE.relative_to( + test_project_paths.TEST_PROJECT_PATH + ).as_posix() + ) + assert f"'basename': '{basename}'" in result + assert f"'relative_path': '{relative_path}'" in result + assert "Text under header A." in result + + +def test_get_parent_node(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + node_id = 30 + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.get_parent_node(node_id, driver) + + assert "ParentNode" in result + assert "ASTNode" in result + + assert "'start_line': 2" in result + assert "'end_line': 2" in result + assert "'type': 'parameter_list'" in result + assert "'text': '()'" in result + + +def test_get_children_node(setup_neo4j_container): + neo4j_container = setup_neo4j_container + uri = neo4j_container.get_connection_url() + + node_id = 20 + with GraphDatabase.driver(uri, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver: + result = graph_traversal.get_children_node(node_id, driver) + + assert "ChildNode" in result + assert "ASTNode" in result + + assert result.count("Result") == 3 + + assert "'start_line': 3" in result + assert "'end_line': 3" in result + assert "'type': 'string_literal'" in result + assert "'text': '\"Hello world!\"'" in result