From 185f936c279cbd398e0a9e2b72b42a049adf9330 Mon Sep 17 00:00:00 2001 From: Leon Luttenberger Date: Mon, 18 Dec 2023 14:13:43 -0600 Subject: [PATCH] fix neptune tests --- tests/unit/test_neptune.py | 8 ++++---- tests/unit/test_neptune_parsing.py | 28 ++++++++++++++++------------ 2 files changed, 20 insertions(+), 16 deletions(-) diff --git a/tests/unit/test_neptune.py b/tests/unit/test_neptune.py index 28f2581d6..831dcb92a 100644 --- a/tests/unit/test_neptune.py +++ b/tests/unit/test_neptune.py @@ -130,10 +130,10 @@ def test_gremlin_query_vertices(neptune_endpoint, neptune_port) -> Dict[str, Any wr.neptune.execute_gremlin(client, f"g.addV().property(T.id, '{uuid.uuid4()}')") df = wr.neptune.execute_gremlin(client, "g.V().limit(1)") - assert df.shape == (1, 2) + assert df.shape == (1, 3) df = wr.neptune.execute_gremlin(client, "g.V().limit(2)") - assert df.shape == (2, 2) + assert df.shape == (2, 3) def test_gremlin_query_edges(neptune_endpoint, neptune_port) -> Dict[str, Any]: @@ -141,10 +141,10 @@ def test_gremlin_query_edges(neptune_endpoint, neptune_port) -> Dict[str, Any]: wr.neptune.execute_gremlin(client, "g.addE('bar').from(addV('foo')).to(addV('foo'))") df = wr.neptune.execute_gremlin(client, "g.E().limit(1)") - assert df.shape == (1, 4) + assert df.shape == (1, 5) df = wr.neptune.execute_gremlin(client, "g.E().limit(2)") - assert df.shape == (2, 4) + assert df.shape == (2, 5) def test_gremlin_query_no_results(neptune_endpoint, neptune_port) -> Dict[str, Any]: diff --git a/tests/unit/test_neptune_parsing.py b/tests/unit/test_neptune_parsing.py index 1bd33dcd9..fbb6b9e58 100644 --- a/tests/unit/test_neptune_parsing.py +++ b/tests/unit/test_neptune_parsing.py @@ -26,9 +26,10 @@ def test_parse_gremlin_vertex_elements(gremlin_parser): out = gremlin_parser.gremlin_results_to_dict(input) df = pd.DataFrame.from_records(out) row = df.iloc[0] - assert df.shape == (1, 2) + assert df.shape == (1, 3) assert row["id"] == "foo" assert row["label"] == "vertex" + assert row["properties"] is None # parse multiple vertex elements v1 = Vertex("bar") @@ -36,9 +37,10 @@ def test_parse_gremlin_vertex_elements(gremlin_parser): out = gremlin_parser.gremlin_results_to_dict(input) df = pd.DataFrame.from_records(out) row = df.iloc[1] - assert df.shape == (2, 2) + assert df.shape == (2, 3) assert row["id"] == "bar" assert row["label"] == "vertex" + assert row["properties"] is None # parse Edge elements @@ -76,22 +78,24 @@ def test_parse_gremlin_property_elements(gremlin_parser): out = gremlin_parser.gremlin_results_to_dict(input) df = pd.DataFrame.from_records(out) row = df.iloc[0] - assert df.shape == (1, 5) + assert df.shape == (1, 6) assert row["id"] == "foo" assert row["label"] == "name" assert row["value"] == "bar" assert row["key"] == "name" assert row["vertex"] == "v1" + assert row["properties"] is None v = Property("foo", "name", "bar") input = [v] out = gremlin_parser.gremlin_results_to_dict(input) df = pd.DataFrame.from_records(out) row = df.iloc[0] - assert df.shape == (1, 3) + assert df.shape == (1, 4) assert row["element"] == "bar" assert row["value"] == "name" assert row["key"] == "foo" + assert row["properties"] is None # parse Path elements @@ -105,9 +109,9 @@ def test_parse_gremlin_path_elements(gremlin_parser): df = pd.DataFrame.from_records(out) row = df.iloc[0] assert df.shape == (1, 3) - assert row[0] == {"id": "foo", "label": "vertex"} - assert row[1] == {"id": "e1", "label": "label", "outV": "foo", "inV": "bar"} - assert row[2] == {"id": "bar", "label": "vertex"} + assert row[0] == {"id": "foo", "label": "vertex", "properties": None} + assert row[1] == {"id": "e1", "label": "label", "outV": "foo", "inV": "bar", "properties": None} + assert row[2] == {"id": "bar", "label": "vertex", "properties": None} # parse path with multiple elements e2 = Edge("bar", "out1", "label", "in2") @@ -117,9 +121,9 @@ def test_parse_gremlin_path_elements(gremlin_parser): df = pd.DataFrame.from_records(out) row = df.iloc[1] assert df.shape == (2, 3) - assert row[0] == {"id": "bar", "label": "vertex"} - assert row[1] == {"id": "bar", "label": "label", "outV": "out1", "inV": "in2"} - assert row[2] == {"id": "in2", "label": "vertex"} + assert row[0] == {"id": "bar", "label": "vertex", "properties": None} + assert row[1] == {"id": "bar", "label": "label", "outV": "out1", "inV": "in2", "properties": None} + assert row[2] == {"id": "in2", "label": "vertex", "properties": None} # parse path with maps p = Path( @@ -129,7 +133,7 @@ def test_parse_gremlin_path_elements(gremlin_parser): out = gremlin_parser.gremlin_results_to_dict([p]) df = pd.DataFrame.from_records(out) row = df.iloc[0] - assert df.shape == (1, 3) + assert df.shape == (1, 4) assert row[0]["name"] == "foo" assert row[0]["age"] == 29 assert row[1]["dist"] == 32 @@ -144,7 +148,7 @@ def test_parse_gremlin_path_elements(gremlin_parser): out = gremlin_parser.gremlin_results_to_dict([p]) df = pd.DataFrame.from_records(out) row = df.iloc[0] - assert df.shape == (1, 3) + assert df.shape == (1, 4) assert row[0]["name"] == "foo" assert row[0]["age"] == 29 assert row[1] == {"id": "bar", "label": "label", "outV": "out1", "inV": "in2"}