From 6764095ed1183a02e451ab03f8908c7cd45209ba Mon Sep 17 00:00:00 2001
From: wonder <hezhaohui0807@163.com>
Date: Mon, 13 Apr 2026 14:45:14 +0800
Subject: [PATCH] =?UTF-8?q?Refactor:=20=E5=B0=81=E8=A3=85=E8=8E=B7?=
 =?UTF-8?q?=E5=8F=96neo4j=E7=9A=84=E8=BF=9E=E6=8E=A5=E5=87=BD=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 backend/neo4j/client.go      | 7 +++----
 backend/neo4j/client_test.go | 6 ++++--
 2 files changed, 7 insertions(+), 6 deletions(-)

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())
 }