Last active
October 17, 2023 19:01
-
-
Save jarifibrahim/6f800a33079722653745ca6fd2f96772 to your computer and use it in GitHub Desktop.
Cleaning up test fixtures with GORM hooks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package main | |
import ( | |
"database/sql" | |
"fmt" | |
"testing" | |
"github.com/jinzhu/gorm" | |
_ "github.com/jinzhu/gorm/dialects/sqlite" | |
) | |
type Product struct { | |
gorm.Model | |
Code string | |
Price uint | |
} | |
// Error checking intentionally skipped at some places | |
func TestCleanup(t *testing.T) { | |
db, _ := gorm.Open("sqlite3", "test.db") | |
defer db.Close() | |
// Setup the cleaner | |
cleaner := DeleteCreatedEntities(db) | |
defer cleaner() | |
// Create | |
db.Create(&Product{Code: "foo", Price: 100}) | |
db.Create(&Product{Code: "bar", Price: 2000}) | |
db.Create(&Product{Code: "fooBar", Price: 400}) | |
// Perform some tests here | |
// ... | |
// ... | |
} | |
func DeleteCreatedEntities(db *gorm.DB) func() { | |
type entity struct { | |
table string | |
keyname string | |
key interface{} | |
} | |
var entries []entity | |
hookName := "cleanupHook" | |
// Setup the onCreate Hook | |
db.Callback().Create().After("gorm:create").Register(hookName, func(scope *gorm.Scope) { | |
fmt.Printf("Inserted entities of %s with %s=%v\n", scope.TableName(), scope.PrimaryKey(), scope.PrimaryKeyValue()) | |
entries = append(entries, entity{table: scope.TableName(), keyname: scope.PrimaryKey(), key: scope.PrimaryKeyValue()}) | |
}) | |
return func() { | |
// Remove the hook once we're done | |
defer db.Callback().Create().Remove(hookName) | |
// Find out if the current db object is already a transaction | |
_, inTransaction := db.CommonDB().(*sql.Tx) | |
tx := db | |
if !inTransaction { | |
tx = db.Begin() | |
} | |
// Loop from the end. It is important that we delete the entries in the | |
// reverse order of their insertion | |
for i := len(entries) - 1; i >= 0; i-- { | |
entry := entries[i] | |
fmt.Printf("Deleting entities from '%s' table with key %v\n", entry.table, entry.key) | |
tx.Table(entry.table).Where(entry.keyname+" = ?", entry.key).Delete("") | |
} | |
if !inTransaction { | |
tx.Commit() | |
} | |
} | |
} |
Gist supporting the latest gorm version here: https://gist.github.com/lundha/d592bd11f1c490999e3bcdb272e114bc
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
This method above doesn't support the latest version.
I also found rolling back the transaction to be easier in general.
https://gorm.io/docs/transactions.html#SavePoint-RollbackTo