another test

tests/unit/test_schema.py

@@ -12,7 +12,44 @@
 #  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 neo4j_genai.schema import get_schema, NODE_PROPERTIES_QUERY, REL_PROPERTIES_QUERY, REL_QUERY, EXCLUDED_LABELS, BASE_ENTITY_LABEL, EXCLUDED_RELS
+from unittest.mock import patch
+from neo4j_genai.schema import (
+    get_schema,
+    NODE_PROPERTIES_QUERY,
+    REL_PROPERTIES_QUERY,
+    REL_QUERY,
+    EXCLUDED_LABELS,
+    BASE_ENTITY_LABEL,
+    EXCLUDED_RELS,
+)
+
+
+def _query_return_value(*args, **kwargs):
+    if NODE_PROPERTIES_QUERY in args[1]:
+        return [
+            {
+                "output": {
+                    "properties": [{"property": "property_a", "type": "STRING"}],
+                    "labels": "LabelA",
+                }
+            }
+        ]
+    if REL_PROPERTIES_QUERY in args[1]:
+        return [
+            {
+                "output": {
+                    "type": "REL_TYPE",
+                    "properties": [{"property": "rel_prop", "type": "STRING"}],
+                }
+            }
+        ]
+    if REL_QUERY in args[1]:
+        return [
+            {"output": {"start": "LabelA", "type": "REL_TYPE", "end": "LabelB"}},
+            {"output": {"start": "LabelA", "type": "REL_TYPE", "end": "LabelC"}},
+        ]
+
+    raise AssertionError("Unexpected query")
 
 
 def test_get_schema_happy_path(driver):
@@ -30,3 +67,18 @@ def test_get_schema_happy_path(driver):
         REL_QUERY,
         {"EXCLUDED_LABELS": EXCLUDED_LABELS + [BASE_ENTITY_LABEL]},
     )
+
+
+@patch("neo4j_genai.schema._query", side_effect=_query_return_value)
+def test_get_schema_ensure_formatted_response(driver):
+    result = get_schema(driver)
+    assert (
+        result
+        == """Node properties:
+LabelA {property_a: STRING}
+Relationship properties:
+REL_TYPE {rel_prop: STRING}
+The relationships:
+(:LabelA)-[:REL_TYPE]->(:LabelB)
+(:LabelA)-[:REL_TYPE]->(:LabelC)"""
+    )