From 510473a09ce058d9315e0d242f93b87de84ebdb3 Mon Sep 17 00:00:00 2001
From: Andrew Glaude <andrew.glaude@datadoghq.com>
Date: Tue, 5 Dec 2023 15:06:32 -0500
Subject: [PATCH] db contribs: Can not defer in TestMain due to os.Exit

---
 contrib/database/sql/sql_test.go             | 6 ++++--
 contrib/gopkg.in/jinzhu/gorm.v1/gorm_test.go | 6 ++++--
 contrib/gorm.io/gorm.v1/gorm_test.go         | 6 ++++--
 contrib/jinzhu/gorm/gorm_test.go             | 6 ++++--
 contrib/jmoiron/sqlx/sql_test.go             | 6 ++++--
 5 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/contrib/database/sql/sql_test.go b/contrib/database/sql/sql_test.go
index c92e2517b1..686dbb6240 100644
--- a/contrib/database/sql/sql_test.go
+++ b/contrib/database/sql/sql_test.go
@@ -38,8 +38,10 @@ func TestMain(m *testing.M) {
 		fmt.Println("--- SKIP: to enable integration test, set the INTEGRATION environment variable")
 		os.Exit(0)
 	}
-	defer sqltest.Prepare(tableName)()
-	os.Exit(m.Run())
+	cleanup := sqltest.Prepare(tableName)
+	testResult := m.Run()
+	cleanup()
+	os.Exit(testResult)
 }
 
 func TestSqlServer(t *testing.T) {
diff --git a/contrib/gopkg.in/jinzhu/gorm.v1/gorm_test.go b/contrib/gopkg.in/jinzhu/gorm.v1/gorm_test.go
index cd0350f78c..91b24eba01 100644
--- a/contrib/gopkg.in/jinzhu/gorm.v1/gorm_test.go
+++ b/contrib/gopkg.in/jinzhu/gorm.v1/gorm_test.go
@@ -35,8 +35,10 @@ func TestMain(m *testing.M) {
 		fmt.Println("--- SKIP: to enable integration test, set the INTEGRATION environment variable")
 		os.Exit(0)
 	}
-	defer sqltest.Prepare(tableName)()
-	os.Exit(m.Run())
+	cleanup := sqltest.Prepare(tableName)
+	testResult := m.Run()
+	cleanup()
+	os.Exit(testResult)
 }
 
 func TestMySQL(t *testing.T) {
diff --git a/contrib/gorm.io/gorm.v1/gorm_test.go b/contrib/gorm.io/gorm.v1/gorm_test.go
index 0b3a9ea28a..46263512a8 100644
--- a/contrib/gorm.io/gorm.v1/gorm_test.go
+++ b/contrib/gorm.io/gorm.v1/gorm_test.go
@@ -44,8 +44,10 @@ func TestMain(m *testing.M) {
 		fmt.Println("--- SKIP: to enable integration test, set the INTEGRATION environment variable")
 		os.Exit(0)
 	}
-	defer sqltest.Prepare(tableName)()
-	os.Exit(m.Run())
+	cleanup := sqltest.Prepare(tableName)
+	testResult := m.Run()
+	cleanup()
+	os.Exit(testResult)
 }
 
 func TestMySQL(t *testing.T) {
diff --git a/contrib/jinzhu/gorm/gorm_test.go b/contrib/jinzhu/gorm/gorm_test.go
index 2c7b07f812..cfaf091185 100644
--- a/contrib/jinzhu/gorm/gorm_test.go
+++ b/contrib/jinzhu/gorm/gorm_test.go
@@ -35,8 +35,10 @@ func TestMain(m *testing.M) {
 		fmt.Println("--- SKIP: to enable integration test, set the INTEGRATION environment variable")
 		os.Exit(0)
 	}
-	defer sqltest.Prepare(tableName)()
-	os.Exit(m.Run())
+	cleanup := sqltest.Prepare(tableName)
+	testResult := m.Run()
+	cleanup()
+	os.Exit(testResult)
 }
 
 func TestSqlServer(t *testing.T) {
diff --git a/contrib/jmoiron/sqlx/sql_test.go b/contrib/jmoiron/sqlx/sql_test.go
index 54322f839a..8b1e7d0ce2 100644
--- a/contrib/jmoiron/sqlx/sql_test.go
+++ b/contrib/jmoiron/sqlx/sql_test.go
@@ -29,8 +29,10 @@ func TestMain(m *testing.M) {
 		fmt.Println("--- SKIP: to enable integration test, set the INTEGRATION environment variable")
 		os.Exit(0)
 	}
-	defer sqltest.Prepare(tableName)()
-	os.Exit(m.Run())
+	cleanup := sqltest.Prepare(tableName)
+	testResult := m.Run()
+	cleanup()
+	os.Exit(testResult)
 }
 
 func TestMySQL(t *testing.T) {