diff --git a/backend/neo4j/client.go b/backend/neo4j/client.go
index b48f5a7..159be3f 100644
--- a/backend/neo4j/client.go
+++ b/backend/neo4j/client.go
@@ -2,14 +2,13 @@ package neo4j
 
 import (
 	"context"
-	"fmt"
 
 	"github.com/neo4j/neo4j-go-driver/v5/neo4j"
 
 	"knowledge-graph-backend/config"
 )
 
-func DoConnect(cfg config.Neo4jConfig) {
+func NewDriver(cfg config.Neo4jConfig) neo4j.DriverWithContext {
 	ctx := context.Background()
 	driver, err := neo4j.NewDriverWithContext(
 		cfg.URI,
@@ -17,11 +16,11 @@ func DoConnect(cfg config.Neo4jConfig) {
 	if err != nil {
 		panic(err)
 	}
-	defer driver.Close(ctx)
 
 	err = driver.VerifyConnectivity(ctx)
 	if err != nil {
 		panic(err)
 	}
-	fmt.Println("Connection established.")
+
+	return driver
 }
diff --git a/backend/neo4j/client_test.go b/backend/neo4j/client_test.go
index 4f91296..f81d2f0 100644
--- a/backend/neo4j/client_test.go
+++ b/backend/neo4j/client_test.go
@@ -1,12 +1,13 @@
 package neo4j
 
 import (
+	"context"
 	"testing"
 
 	"knowledge-graph-backend/config"
 )
 
-func TestDoConnect(t *testing.T) {
+func TestNewDriver(t *testing.T) {
 	cnf, err := config.Load()
 	if err != nil {
 		t.Errorf("[cnf]: %v", err)
@@ -16,5 +17,6 @@ func TestDoConnect(t *testing.T) {
 		Username: cnf.Neo4j.Username,
 		Password: cnf.Neo4j.Password,
 	}
-	DoConnect(cfg)
+	driver := NewDriver(cfg)
+	defer driver.Close(context.Background())
 }