test: refactor mongo connect helper function

reearth · May 11, 2022 · 751e666 · 751e666
1 parent cb82d05
commit 751e666
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/internal/infrastructure/mongo/asset_test.go b/internal/infrastructure/mongo/asset_test.go
@@ -36,16 +36,15 @@ func TestFindByID(t *testing.T) {
 		},
 	}
 
-	initDB := connect(t)
+	init := connect(t)
 
 	for _, tc := range tests {
 		tc := tc
 
 		t.Run(tc.Name, func(t *testing.T) {
 			t.Parallel()
 
-			client, dropDB := initDB()
-			defer dropDB()
+			client := init(t)
 
 			repo := NewAsset(client)
 			ctx := context.Background()

diff --git a/internal/infrastructure/mongo/mongo_test.go b/internal/infrastructure/mongo/mongo_test.go
@@ -13,7 +13,7 @@ import (
 	"go.mongodb.org/mongo-driver/x/mongo/driver/uuid"
 )
 
-func connect(t *testing.T) func() (*mongodoc.Client, func()) {
+func connect(t *testing.T) func(*testing.T) *mongodoc.Client {
 	t.Helper()
 
 	// Skip unit testing if "REEARTH_DB" is not configured
@@ -31,13 +31,17 @@ func connect(t *testing.T) func() (*mongodoc.Client, func()) {
 			SetConnectTimeout(time.Second*10),
 	)
 
-	return func() (*mongodoc.Client, func()) {
+	return func(t *testing.T) *mongodoc.Client {
+		t.Helper()
+
 		database, _ := uuid.New()
 		databaseName := "reearth-test-" + hex.EncodeToString(database[:])
 		client := mongodoc.NewClient(databaseName, c)
 
-		return client, func() {
+		t.Cleanup(func() {
 			_ = c.Database(databaseName).Drop(context.Background())
-		}
+		})
+
+		return client
 	}
 }