fix neptune tests

tests/unit/test_neptune.py

@@ -130,21 +130,21 @@ 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]:
     client = wr.neptune.connect(neptune_endpoint, neptune_port, iam_enabled=False)
 
     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]:

tests/unit/test_neptune_parsing.py

@@ -26,19 +26,21 @@ 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")
     input = [v, v1]
     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"}