From 7d183741976645c8d5945d24f6f7cb8b1f54f4b7 Mon Sep 17 00:00:00 2001
From: Eric Solender <ericsolender@gmail.com>
Date: Wed, 18 Sep 2019 12:41:38 -0400
Subject: [PATCH] fixed missing WithNeo calls

---
 decoder_test.go | 2 +-
 delete_test.go  | 2 +-
 save.go         | 6 ++++--
 save_test.go    | 2 +-
 4 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/decoder_test.go b/decoder_test.go
index 152ab70..8e49cd8 100644
--- a/decoder_test.go
+++ b/decoder_test.go
@@ -38,7 +38,7 @@ func TestDecode(t *testing.T){
 	if err != nil {
 		require.Nil(t, err)
 	}
-	driverPool.Reclaim(conn)
+	defer driverPool.Reclaim(conn)
 
 
 	rows, err := dsl.QB().WithNeo(conn).Cypher(query).Query(nil)
diff --git a/delete_test.go b/delete_test.go
index d4e4924..2b45321 100644
--- a/delete_test.go
+++ b/delete_test.go
@@ -15,7 +15,7 @@ func TestDelete(t *testing.T){
 	if err != nil {
 		require.Nil(t, err)
 	}
-	driverPool.Reclaim(conn)
+	defer driverPool.Reclaim(conn)
 
 	del := a{
 		Id: 0,
diff --git a/save.go b/save.go
index 496de59..8cca171 100644
--- a/save.go
+++ b/save.go
@@ -81,7 +81,7 @@ func saveDepth(sess *driver.BoltConn, obj interface{}, depth int) error {
 	return relateNodes(sess, relations, ids)
 }
 
-func createNodes(sess *driver.BoltConn, crNodes map[string]map[string]nodeCreateConf) (map[string]int64, error){
+func createNodes(conn *driver.BoltConn, crNodes map[string]map[string]nodeCreateConf) (map[string]int64, error){
 	idMap := map[string]int64{}
 
 	for label, nodes := range crNodes{
@@ -124,6 +124,7 @@ func createNodes(sess *driver.BoltConn, crNodes map[string]map[string]nodeCreate
 				},
 				Alias: "id",
 			}).
+			WithNeo(conn).
 			Query(map[string]interface{}{
 				"rows": rows,
 			})
@@ -157,7 +158,7 @@ func createNodes(sess *driver.BoltConn, crNodes map[string]map[string]nodeCreate
 	return idMap, nil
 }
 
-func relateNodes(sess *driver.BoltConn, relations map[string][]relCreateConf, ids map[string]int64) error{
+func relateNodes(conn *driver.BoltConn, relations map[string][]relCreateConf, ids map[string]int64) error{
 	if relations == nil || len(relations) == 0{
 		return errors.New("relations can not be nil or empty")
 	}
@@ -245,6 +246,7 @@ func relateNodes(sess *driver.BoltConn, relations map[string][]relCreateConf, id
 				Path: mergePath,
 			}).
 			Cypher("SET rel += row.props").
+			WithNeo(conn).
 			Exec(map[string]interface{}{
 				"rows": params,
 			})
diff --git a/save_test.go b/save_test.go
index fcbbf93..9b72ae3 100644
--- a/save_test.go
+++ b/save_test.go
@@ -150,7 +150,7 @@ func TestSave(t *testing.T){
 	if err != nil {
 		require.Nil(t, err)
 	}
-	driverPool.Reclaim(conn)
+	defer driverPool.Reclaim(conn)
 
 	req.Nil(saveDepth(conn, comp2, defaultSaveDepth))
 }
\ No newline at end of file