feat: add db tests

.github/workflows/tests.yml

@@ -24,6 +24,15 @@ jobs:
           - 6379:6379
         options: -e REDIS_PASSWORD=difyai123456
 
+      postgres:
+        image: postgres
+        env:
+          POSTGRES_USER: postgres
+          POSTGRES_PASSWORD: difyai123456
+          POSTGRES_DB: difyai
+        ports:
+          - 5432:5432
+
     steps:
       - uses: actions/checkout@v2
 

internal/db/init.go

@@ -103,3 +103,19 @@ func Init(config *app.Config) {
 
 	log.Info("dify plugin db initialized")
 }
+
+func Close() {
+	db, err := DifyPluginDB.DB()
+	if err != nil {
+		log.Error("failed to close dify plugin db: %v", err)
+		return
+	}
+
+	err = db.Close()
+	if err != nil {
+		log.Error("failed to close dify plugin db: %v", err)
+		return
+	}
+
+	log.Info("dify plugin db closed")
+}

internal/db/pg_test.go

@@ -0,0 +1,50 @@
+package db
+
+import (
+	"errors"
+	"testing"
+
+	"gorm.io/gorm"
+)
+
+func TestTransaction(t *testing.T) {
+	if err := initDifyPluginDB("0.0.0.0", 5432, "testing", "postgres", "difyai123456", "disable"); err != nil {
+		t.Fatal(err)
+	}
+	defer Close()
+
+	type TestTable struct {
+		gorm.Model
+	}
+
+	err := CreateTable(&TestTable{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer DropTable(&TestTable{})
+
+	err = WithTransaction(func(tx *gorm.DB) error {
+		data := TestTable{}
+		err := tx.Create(&data).Error
+		if err != nil {
+			return err
+		}
+
+		return errors.New("rollback")
+	})
+
+	if err == nil {
+		t.Fatal("expected error")
+	} else if err.Error() != "rollback" {
+		t.Fatal("unexpected error")
+	}
+
+	count, err := GetCount[TestTable]()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if count != 0 {
+		t.Fatal("unexpected count")
+	}
+}

internal/db/pgsql.go

@@ -311,11 +311,6 @@ func DropTable(model any) error {
 }
 
 // NOTE: not used in production, only for testing
-func DropDatabase(dbname string) error {
-	return DifyPluginDB.Exec(fmt.Sprintf("DROP DATABASE %s", dbname)).Error
-}
-
-// NOTE: not used in production, only for testing
 func CreateDatabase(dbname string) error {
 	return DifyPluginDB.Exec(fmt.Sprintf("CREATE DATABASE %s", dbname)).Error
 }