참고소스 수정본

참고/neo4j-graphrag-python-main/tests/e2e/test_schema_e2e.py

@@ -0,0 +1,222 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from __future__ import annotations
+
+import pytest
+from neo4j import Driver
+from neo4j_graphrag.schema import (
+    BASE_ENTITY_LABEL,
+    NODE_PROPERTIES_QUERY,
+    REL_PROPERTIES_QUERY,
+    REL_QUERY,
+    get_structured_schema,
+    query_database,
+)
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_cypher_returns_correct_node_properties(driver: Driver) -> None:
+    node_properties = query_database(
+        driver,
+        NODE_PROPERTIES_QUERY,
+        params={"EXCLUDED_LABELS": [BASE_ENTITY_LABEL], "SAMPLE": 1000},
+    )
+
+    expected_node_properties = [
+        {
+            "output": {
+                "properties": [{"property": "property_a", "type": "STRING"}],
+                "label": "LabelA",
+            }
+        }
+    ]
+
+    assert node_properties == expected_node_properties
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_cypher_returns_correct_relationship_properties(driver: Driver) -> None:
+    relationships_properties = query_database(
+        driver,
+        REL_PROPERTIES_QUERY,
+        params={"EXCLUDED_LABELS": [BASE_ENTITY_LABEL], "SAMPLE": 1000},
+    )
+
+    expected_relationships_properties = [
+        {
+            "output": {
+                "type": "REL_TYPE",
+                "properties": [{"property": "rel_prop", "type": "STRING"}],
+            }
+        }
+    ]
+
+    assert relationships_properties == expected_relationships_properties
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_cypher_returns_correct_relationships(driver: Driver) -> None:
+    relationships = query_database(
+        driver,
+        REL_QUERY,
+        params={"EXCLUDED_LABELS": [BASE_ENTITY_LABEL], "SAMPLE": 1000},
+    )
+
+    expected_relationships = [
+        {"output": {"start": "LabelA", "type": "REL_TYPE", "end": "LabelB"}},
+        {"output": {"start": "LabelA", "type": "REL_TYPE", "end": "LabelC"}},
+    ]
+
+    assert (
+        sorted(relationships, key=lambda x: x["output"]["end"])
+        == expected_relationships
+    )
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_get_structured_schema_returns_correct_node_properties(driver: Driver) -> None:
+    result = get_structured_schema(driver)
+    assert result["node_props"]["LabelA"] == [
+        {"property": "property_a", "type": "STRING"}
+    ]
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_get_enhanced_structured_schema_returns_correct_node_properties(
+    driver: Driver,
+) -> None:
+    result = get_structured_schema(driver, True)
+    assert result["node_props"]["LabelA"] == [
+        {
+            "property": "property_a",
+            "type": "STRING",
+            "values": ["a"],
+            "distinct_count": 1,
+        }
+    ]
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_get_structured_schema_returns_correct_relationship_properties(
+    driver: Driver,
+) -> None:
+    result = get_structured_schema(driver)
+    assert result["rel_props"]["REL_TYPE"] == [
+        {"property": "rel_prop", "type": "STRING"}
+    ]
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_get_enhanced_structured_schema_returns_correct_relationship_properties(
+    driver: Driver,
+) -> None:
+    result = get_structured_schema(driver, True)
+    assert result["rel_props"]["REL_TYPE"] == [
+        {
+            "property": "rel_prop",
+            "type": "STRING",
+            "values": ["abc"],
+            "distinct_count": 1,
+        }
+    ]
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_get_structured_schema_returns_correct_relationships(driver: Driver) -> None:
+    result = get_structured_schema(driver)
+    assert sorted(result["relationships"], key=lambda x: x["end"]) == [
+        {"end": "LabelB", "start": "LabelA", "type": "REL_TYPE"},
+        {"end": "LabelC", "start": "LabelA", "type": "REL_TYPE"},
+    ]
+
+
+@pytest.mark.enterprise_only
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_get_structured_schema_returns_correct_constraints(driver: Driver) -> None:
+    query_database(driver, "DROP CONSTRAINT test_constraint IF EXISTS")
+    query_database(
+        driver,
+        "CREATE CONSTRAINT test_constraint IF NOT EXISTS FOR (n:LabelA) REQUIRE n.property_a IS NOT NULL",
+    )
+
+    result = get_structured_schema(driver)
+    assert result["metadata"]["constraint"][0].get("name") == "test_constraint"
+    assert result["metadata"]["constraint"][0].get("type") == "NODE_PROPERTY_EXISTENCE"
+    assert result["metadata"]["constraint"][0].get("entityType") == "NODE"
+    assert result["metadata"]["constraint"][0].get("labelsOrTypes") == ["LabelA"]
+    assert result["metadata"]["constraint"][0].get("properties") == ["property_a"]
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_get_structured_schema_returns_correct_indexes(driver: Driver) -> None:
+    query_database(driver, "DROP INDEX node_range_index IF EXISTS")
+    query_database(
+        driver, "CREATE INDEX node_range_index FOR (n:LabelA) ON (n.property_a)"
+    )
+
+    result = get_structured_schema(driver)
+    assert result["metadata"]["index"][0].get("label") == "LabelA"
+    assert result["metadata"]["index"][0].get("properties") == ["property_a"]
+    assert result["metadata"]["index"][0].get("type") == "RANGE"
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_neo4j_sanitize_values(driver: Driver) -> None:
+    output = query_database(driver, "RETURN range(0,130,1) AS result", sanitize=True)
+    assert output == [{}]
+
+
+def test_enhanced_schema_exception(driver: Driver) -> None:
+    query_database(driver, "MATCH (n) DETACH DELETE n")
+    query_database(
+        driver,
+        "CREATE (:Node {foo: 'bar'}), (:Node {foo: 1}), (:Node {foo: [1,2]}), "
+        "(: EmptyNode)",
+    )
+    query_database(
+        driver,
+        "MATCH (a:Node {foo: 'bar'}), (b:Node {foo: 1}), "
+        "(c:Node {foo: [1,2]}), (d: EmptyNode) "
+        "CREATE (a)-[:REL {foo: 'bar'}]->(b), (b)-[:REL {foo: 1}]->(c), "
+        "(c)-[:REL {foo: [1,2]}]->(a), (d)-[:EMPTY_REL {}]->(d)",
+    )
+    result = get_structured_schema(driver, True)
+    del result["metadata"]
+
+    assert list(result.keys()) == ["node_props", "rel_props", "relationships"]
+    node_props = result["node_props"]
+    assert list(node_props.keys()) == ["Node"]
+    assert len(node_props["Node"]) == 1
+    assert list(node_props["Node"][0].keys()) == ["property", "type"]
+    assert node_props["Node"][0]["property"] == "foo"
+    assert node_props["Node"][0]["type"] in ["STRING", "INTEGER", "LIST"]
+
+    rel_props = result["rel_props"]
+    assert list(rel_props.keys()) == ["REL"]
+    assert len(rel_props["REL"]) == 1
+    assert list(rel_props["REL"][0].keys()) == ["property", "type"]
+    assert rel_props["REL"][0]["property"] == "foo"
+    assert rel_props["REL"][0]["type"] in ["STRING", "INTEGER", "LIST"]
+
+    expected_rels = [
+        {
+            "end": "Node",
+            "start": "Node",
+            "type": "REL",
+        },
+        {"end": "EmptyNode", "start": "EmptyNode", "type": "EMPTY_REL"},
+    ]
+    rels = result["relationships"]
+    assert rels == expected_rels