[usage] Automatically cleanup records between test DB connections

Author: easyCZ

components/usage/pkg/db/conn.go

@@ -72,12 +72,26 @@ func ConnectForTests(t *testing.T) *gorm.DB {
 	})
 	require.NoError(t, err, "Failed to establish connection to DB. In a workspace, run `leeway build components/usage:init-testdb` once to bootstrap the DB.")
 
+	conn = conn.Debug()
+
 	t.Cleanup(func() {
+		// Delete records for known models from the DB
+		log.Info("Cleaning up DB between connections.")
+		for _, model := range []interface{}{
+			&WorkspaceInstance{},
+			&Workspace{},
+			&Project{},
+			&Team{},
+		} {
+			// See https://gorm.io/docs/delete.html#Block-Global-Delete
+			tx := conn.Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(model)
+			require.NoError(t, tx.Error)
+		}
+
 		rawConn, err := conn.DB()
 		require.NoError(t, err)
-
 		require.NoError(t, rawConn.Close(), "must close database connection")
 	})
 
-	return conn.Debug()
+	return conn
 }

components/usage/pkg/db/project_test.go

@@ -64,10 +64,5 @@ func insertRawProject(t *testing.T, conn *gorm.DB, obj map[string]interface{}) u
 	tx := conn.Exec(statement, values)
 	require.NoError(t, tx.Error)
 
-	t.Cleanup(func() {
-		tx = conn.Delete(&db.Project{ID: id})
-		require.NoError(t, tx.Error)
-	})
-
 	return id
 }

components/usage/pkg/db/workspace_instance_test.go

@@ -70,11 +70,6 @@ func insertRawWorkspaceInstance(t *testing.T, conn *gorm.DB, object map[string]i
 	statement := fmt.Sprintf(`INSERT INTO d_b_workspace_instance (%s) VALUES ?;`, strings.Join(columns, ", "))
 	id := uuid.MustParse(insertRawObject(t, conn, columns, statement, object))
 
-	t.Cleanup(func() {
-		tx := conn.Delete(&db.WorkspaceInstance{ID: id})
-		require.NoError(t, tx.Error)
-	})
-
 	return id
 }
 

components/usage/pkg/db/workspace_test.go

@@ -73,10 +73,6 @@ func insertRawWorkspace(t *testing.T, conn *gorm.DB, rawWorkspace map[string]int
 	) VALUES ?;`, strings.Join(columns, ", "))
 	id := insertRawObject(t, conn, columns, statement, rawWorkspace)
 
-	t.Cleanup(func() {
-		tx := conn.Delete(&db.Workspace{ID: id})
-		require.NoError(t, tx.Error)
-	})
 	return id
 }