refactor db context helper to consider unique database names

MGTheTrain · Nov 21, 2024 · e2adaff · e2adaff
1 parent 7da7a1c
commit e2adaff
Show file tree

Hide file tree

Showing 5 changed files with 78 additions and 98 deletions.
diff --git a/test/helpers/test_db_context.go b/test/helpers/test_db_context.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"testing"
 
+	"github.com/google/uuid"
 	"gorm.io/driver/postgres"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
@@ -37,37 +38,33 @@ func SetupTestDB(t *testing.T) *TestDBContext {
 			t.Fatalf("POSTGRES_DSN environment variable is not set")
 		}
 
+		// Connect to PostgreSQL without specifying a database (so we can create one if necessary)
 		db, err = gorm.Open(postgres.Open(dsn), &gorm.Config{})
 		if err != nil {
 			t.Fatalf("Failed to connect to PostgreSQL: %v", err)
 		}
 
-		// Ensure the `blobs` database exists
+		// Generate a unique database name using UUID
+		uniqueDBName := "blobs_" + uuid.New().String()
+
+		// Ensure the unique `blobs` database exists, create if necessary
 		sqlDB, err := db.DB()
 		if err != nil {
 			t.Fatalf("Failed to get raw DB connection: %v", err)
 		}
 
-		// Query to check if the `blobs` database exists
-		var dbExists bool
-		err = sqlDB.QueryRow("SELECT 1 FROM pg_database WHERE datname = 'blobs'").Scan(&dbExists)
+		// Create the new database
+		_, err = sqlDB.Exec(fmt.Sprintf("CREATE DATABASE %s", uniqueDBName))
 		if err != nil {
-			if err.Error() != "sql: no rows in result set" {
-				t.Fatalf("Failed to check if database exists: %v", err)
-			}
-			// The database does not exist, create it
-			_, err = sqlDB.Exec("CREATE DATABASE blobs")
-			if err != nil {
-				t.Fatalf("Failed to create database: %v", err)
-			}
-			fmt.Println("Database 'blobs' created successfully.")
+			t.Fatalf("Failed to create database '%s': %v", uniqueDBName, err)
 		}
+		fmt.Printf("Database '%s' created successfully.\n", uniqueDBName)
 
-		// Open the connection to `blobs` database
-		dsn = "user=postgres password=postgres host=localhost port=5432 dbname=blobs sslmode=disable"
+		// Now that the unique `blobs` database is created, connect to it
+		dsn = fmt.Sprintf("user=postgres password=postgres host=localhost port=5432 dbname=%s sslmode=disable", uniqueDBName)
 		db, err = gorm.Open(postgres.Open(dsn), &gorm.Config{})
 		if err != nil {
-			t.Fatalf("Failed to connect to PostgreSQL database 'blobs': %v", err)
+			t.Fatalf("Failed to connect to PostgreSQL database '%s': %v", uniqueDBName, err)
 		}
 
 	case "sqlite":
@@ -97,11 +94,36 @@ func SetupTestDB(t *testing.T) *TestDBContext {
 	}
 }
 
-// TeardownTestDB closes the DB connection after the test
-func TeardownTestDB(t *testing.T, ctx *TestDBContext) {
+// TeardownTestDB closes the DB connection and cleans up the database after the test
+func TeardownTestDB(t *testing.T, ctx *TestDBContext, dbType string) {
 	sqlDB, err := ctx.DB.DB()
 	if err != nil {
 		t.Fatalf("Failed to get DB connection: %v", err)
 	}
-	sqlDB.Close()
+
+	// If using PostgreSQL, drop the unique database created during the test
+	if dbType == "postgres" {
+		// Get the database name from the DSN or context (you might store it during DB setup)
+		databaseName := ctx.DB.Migrator().CurrentDatabase()
+
+		// Close the current DB connection before dropping the database
+		sqlDB.Close()
+
+		// Connect again to PostgreSQL without specifying a database (connect to the default one)
+		dsn := "user=postgres password=postgres host=localhost port=5432 sslmode=disable"
+		db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
+		if err != nil {
+			t.Fatalf("Failed to reconnect to PostgreSQL: %v", err)
+		}
+
+		// Drop the unique database
+		tx := db.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s", databaseName))
+		if tx.Error != nil {
+			t.Fatalf("Failed to drop database '%s': %v", databaseName, tx.Error)
+		}
+		fmt.Printf("Database '%s' dropped successfully.\n", databaseName)
+	} else {
+		// For SQLite, no need to drop the in-memory database, just close the connection
+		sqlDB.Close()
+	}
 }
diff --git a/test/integration/app/services/key_services_test.go b/test/integration/app/services/key_services_test.go
@@ -35,13 +35,10 @@ func removeTestFile(testFilePath string) error {
 
 // Test case for successful file upload and metadata creation
 func TestCryptoKeyUploadService_Upload_Success(t *testing.T) {
-	// Set up environment variable
-	err := os.Setenv("DB_TYPE", "postgres")
-	require.NoError(t, err, "Error setting environment variable")
-
-	// Set up test context and ensure proper teardown
+	// Set up test context
 	ctx := helpers.SetupTestDB(t)
-	defer helpers.TeardownTestDB(t, ctx)
+	dbType := "sqlite"
+	defer helpers.TeardownTestDB(t, ctx, dbType)
 
 	// Initialize Vault Connector
 	connectionString := "DefaultEndpointsProtocol=http;AccountName=devstoreaccount1;AccountKey=Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==;BlobEndpoint=http://127.0.0.1:10000/devstoreaccount1;"
@@ -82,13 +79,10 @@ func TestCryptoKeyUploadService_Upload_Success(t *testing.T) {
 
 // Test case for successful retrieval of cryptographic key metadata by ID
 func TestCryptoKeyMetadataService_GetByID_Success(t *testing.T) {
-	// Set up environment variable
-	err := os.Setenv("DB_TYPE", "postgres")
-	require.NoError(t, err, "Error setting environment variable")
-
-	// Set up test context and ensure proper teardown
+	// Set up test context
 	ctx := helpers.SetupTestDB(t)
-	defer helpers.TeardownTestDB(t, ctx)
+	dbType := "sqlite"
+	defer helpers.TeardownTestDB(t, ctx, dbType)
 
 	// Create a key metadata instance in the database for testing
 	cryptoKeyMeta := &keys.CryptoKeyMeta{
@@ -100,7 +94,7 @@ func TestCryptoKeyMetadataService_GetByID_Success(t *testing.T) {
 	}
 
 	// Store the cryptographic key metadata in the database
-	err = ctx.CryptoKeyRepo.Create(cryptoKeyMeta)
+	err := ctx.CryptoKeyRepo.Create(cryptoKeyMeta)
 	require.NoError(t, err, "Error creating test cryptographic key metadata")
 
 	// Set up the CryptoKeyMetadataService
@@ -123,13 +117,10 @@ func TestCryptoKeyMetadataService_GetByID_Success(t *testing.T) {
 
 // Test case for successful deletion of cryptographic key metadata by ID
 func TestCryptoKeyMetadataService_DeleteByID_Success(t *testing.T) {
-	// Set up environment variable
-	err := os.Setenv("DB_TYPE", "postgres")
-	require.NoError(t, err, "Error setting environment variable")
-
-	// Set up test context and ensure proper teardown
+	// Set up test context
 	ctx := helpers.SetupTestDB(t)
-	defer helpers.TeardownTestDB(t, ctx)
+	dbType := "sqlite"
+	defer helpers.TeardownTestDB(t, ctx, dbType)
 
 	// Create a key metadata instance in the database for testing
 	cryptoKeyMeta := &keys.CryptoKeyMeta{
@@ -141,7 +132,7 @@ func TestCryptoKeyMetadataService_DeleteByID_Success(t *testing.T) {
 	}
 
 	// Store the cryptographic key metadata in the database
-	err = ctx.CryptoKeyRepo.Create(cryptoKeyMeta)
+	err := ctx.CryptoKeyRepo.Create(cryptoKeyMeta)
 	require.NoError(t, err, "Error creating test cryptographic key metadata")
 
 	// Set up the CryptoKeyMetadataService
@@ -164,13 +155,10 @@ func TestCryptoKeyMetadataService_DeleteByID_Success(t *testing.T) {
 
 // Test case for successful download of cryptographic key
 func TestCryptoKeyDownloadService_Download_Success(t *testing.T) {
-	// Set up environment variable
-	err := os.Setenv("DB_TYPE", "postgres")
-	require.NoError(t, err, "Error setting environment variable")
-
-	// Set up test context and ensure proper teardown
+	// Set up test context
 	ctx := helpers.SetupTestDB(t)
-	defer helpers.TeardownTestDB(t, ctx)
+	dbType := "sqlite"
+	defer helpers.TeardownTestDB(t, ctx, dbType)
 
 	// Initialize Vault Connector
 	connectionString := "DefaultEndpointsProtocol=http;AccountName=devstoreaccount1;AccountKey=Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==;BlobEndpoint=http://127.0.0.1:10000/devstoreaccount1;"

diff --git a/test/integration/persistence/repository/psql_blob_repository_test.go b/test/integration/persistence/repository/psql_blob_repository_test.go
@@ -6,7 +6,6 @@ import (
 	"crypto_vault_service/internal/domain/keys"
 	"crypto_vault_service/test/helpers"
 
-	"os"
 	"testing"
 	"time"
 
@@ -15,14 +14,10 @@ import (
 )
 
 func TestBlobPsqlRepository_Create(t *testing.T) {
-	err := os.Setenv("DB_TYPE", "postgres")
-	if err != nil {
-		t.Fatalf("Error setting environment variable: %v", err)
-	}
-
 	// Set up test context
 	ctx := helpers.SetupTestDB(t)
-	defer helpers.TeardownTestDB(t, ctx)
+	dbType := "postgres"
+	defer helpers.TeardownTestDB(t, ctx, dbType)
 
 	// Create a valid CryptoKey object
 	cryptographicKey := keys.CryptoKeyMeta{
@@ -46,7 +41,7 @@ func TestBlobPsqlRepository_Create(t *testing.T) {
 	}
 
 	// Call the Create method
-	err = ctx.BlobRepo.Create(blob)
+	err := ctx.BlobRepo.Create(blob)
 	assert.NoError(t, err, "Create should not return an error")
 
 	// Verify the blob is created and exists in DB
@@ -58,14 +53,10 @@ func TestBlobPsqlRepository_Create(t *testing.T) {
 }
 
 func TestBlobPsqlRepository_GetById(t *testing.T) {
-	err := os.Setenv("DB_TYPE", "postgres")
-	if err != nil {
-		t.Fatalf("Error setting environment variable: %v", err)
-	}
-
 	// Set up test context
 	ctx := helpers.SetupTestDB(t)
-	defer helpers.TeardownTestDB(t, ctx)
+	dbType := "postgres"
+	defer helpers.TeardownTestDB(t, ctx, dbType)
 
 	// Create a valid CryptoKey object
 	cryptographicKey := keys.CryptoKeyMeta{
@@ -89,7 +80,7 @@ func TestBlobPsqlRepository_GetById(t *testing.T) {
 	}
 
 	// Create the blob in DB
-	err = ctx.BlobRepo.Create(blob)
+	err := ctx.BlobRepo.Create(blob)
 	assert.NoError(t, err, "Create should not return an error")
 
 	// Get the blob by ID

diff --git a/test/integration/persistence/repository/sqlite_blob_repository_test.go b/test/integration/persistence/repository/sqlite_blob_repository_test.go
@@ -6,7 +6,6 @@ import (
 	"crypto_vault_service/internal/domain/keys"
 	"crypto_vault_service/test/helpers"
 
-	"os"
 	"testing"
 	"time"
 
@@ -15,14 +14,10 @@ import (
 )
 
 func TestBlobSqliteRepository_Create(t *testing.T) {
-	err := os.Setenv("DB_TYPE", "sqlite")
-	if err != nil {
-		t.Fatalf("Error setting environment variable: %v", err)
-	}
-
 	// Set up test context
 	ctx := helpers.SetupTestDB(t)
-	defer helpers.TeardownTestDB(t, ctx)
+	dbType := "sqlite"
+	defer helpers.TeardownTestDB(t, ctx, dbType)
 
 	// Create a valid CryptoKey object
 	cryptographicKey := keys.CryptoKeyMeta{
@@ -46,7 +41,7 @@ func TestBlobSqliteRepository_Create(t *testing.T) {
 	}
 
 	// Call the Create method
-	err = ctx.BlobRepo.Create(blob)
+	err := ctx.BlobRepo.Create(blob)
 	assert.NoError(t, err, "Create should not return an error")
 
 	// Verify the blob is created and exists in DB
@@ -58,14 +53,10 @@ func TestBlobSqliteRepository_Create(t *testing.T) {
 }
 
 func TestBlobSqliteRepository_GetById(t *testing.T) {
-	err := os.Setenv("DB_TYPE", "sqlite")
-	if err != nil {
-		t.Fatalf("Error setting environment variable: %v", err)
-	}
-
 	// Set up test context
 	ctx := helpers.SetupTestDB(t)
-	defer helpers.TeardownTestDB(t, ctx)
+	dbType := "sqlite"
+	defer helpers.TeardownTestDB(t, ctx, dbType)
 
 	// Create a valid CryptoKey object
 	cryptographicKey := keys.CryptoKeyMeta{
@@ -89,7 +80,7 @@ func TestBlobSqliteRepository_GetById(t *testing.T) {
 	}
 
 	// Create the blob in DB
-	err = ctx.BlobRepo.Create(blob)
+	err := ctx.BlobRepo.Create(blob)
 	assert.NoError(t, err, "Create should not return an error")
 
 	// Get the blob by ID

diff --git a/test/integration/persistence/repository/sqlite_key_repository_test.go b/test/integration/persistence/repository/sqlite_key_repository_test.go
@@ -4,7 +4,6 @@ import (
 	"crypto_vault_service/internal/domain/keys"
 	"crypto_vault_service/test/helpers"
 
-	"os"
 	"testing"
 	"time"
 
@@ -15,14 +14,10 @@ import (
 
 // TestCryptoKeySqliteRepository_Create tests the Create method of GormCryptoKeyRepository
 func TestCryptoKeySqliteRepository_Create(t *testing.T) {
-	err := os.Setenv("DB_TYPE", "sqlite")
-	if err != nil {
-		t.Fatalf("Error setting environment variable: %v", err)
-	}
-
 	// Set up test context
 	ctx := helpers.SetupTestDB(t)
-	defer helpers.TeardownTestDB(t, ctx)
+	dbType := "sqlite"
+	defer helpers.TeardownTestDB(t, ctx, dbType)
 
 	// Create a valid CryptoKey object
 	cryptoKeyMeta := &keys.CryptoKeyMeta{
@@ -34,7 +29,7 @@ func TestCryptoKeySqliteRepository_Create(t *testing.T) {
 	}
 
 	// Create the cryptographic key in DB
-	err = ctx.CryptoKeyRepo.Create(cryptoKeyMeta)
+	err := ctx.CryptoKeyRepo.Create(cryptoKeyMeta)
 	assert.NoError(t, err, "Create should not return an error")
 
 	// Verify the cryptographic key is created and exists in DB
@@ -47,14 +42,10 @@ func TestCryptoKeySqliteRepository_Create(t *testing.T) {
 
 // TestCryptoKeySqliteRepository_GetByID tests the GetByID method of GormCryptoKeyRepository
 func TestCryptoKeySqliteRepository_GetByID(t *testing.T) {
-	err := os.Setenv("DB_TYPE", "sqlite")
-	if err != nil {
-		t.Fatalf("Error setting environment variable: %v", err)
-	}
-
 	// Set up test context
 	ctx := helpers.SetupTestDB(t)
-	defer helpers.TeardownTestDB(t, ctx)
+	dbType := "sqlite"
+	defer helpers.TeardownTestDB(t, ctx, dbType)
 
 	// Create a valid CryptoKey object
 	cryptoKeyMeta := &keys.CryptoKeyMeta{
@@ -66,7 +57,7 @@ func TestCryptoKeySqliteRepository_GetByID(t *testing.T) {
 	}
 
 	// Create the cryptographic key in DB
-	err = ctx.CryptoKeyRepo.Create(cryptoKeyMeta)
+	err := ctx.CryptoKeyRepo.Create(cryptoKeyMeta)
 	assert.NoError(t, err, "Create should not return an error")
 
 	// Get the cryptographic key by ID
@@ -78,14 +69,10 @@ func TestCryptoKeySqliteRepository_GetByID(t *testing.T) {
 
 // TestCryptoKeySqliteRepository_UpdateByID tests the UpdateByID method of GormCryptoKeyRepository
 func TestCryptoKeySqliteRepository_UpdateByID(t *testing.T) {
-	err := os.Setenv("DB_TYPE", "sqlite")
-	if err != nil {
-		t.Fatalf("Error setting environment variable: %v", err)
-	}
-
 	// Set up test context
 	ctx := helpers.SetupTestDB(t)
-	defer helpers.TeardownTestDB(t, ctx)
+	dbType := "sqlite"
+	defer helpers.TeardownTestDB(t, ctx, dbType)
 
 	// Create a valid CryptoKey object
 	cryptoKeyMeta := &keys.CryptoKeyMeta{
@@ -97,7 +84,7 @@ func TestCryptoKeySqliteRepository_UpdateByID(t *testing.T) {
 	}
 
 	// Create the cryptographic key in DB
-	err = ctx.CryptoKeyRepo.Create(cryptoKeyMeta)
+	err := ctx.CryptoKeyRepo.Create(cryptoKeyMeta)
 	assert.NoError(t, err, "Create should not return an error")
 
 	// Update the cryptographic key's type
@@ -116,7 +103,8 @@ func TestCryptoKeySqliteRepository_UpdateByID(t *testing.T) {
 func TestCryptoKeySqliteRepository_DeleteByID(t *testing.T) {
 	// Set up test context
 	ctx := helpers.SetupTestDB(t)
-	defer helpers.TeardownTestDB(t, ctx)
+	dbType := "sqlite"
+	defer helpers.TeardownTestDB(t, ctx, dbType)
 
 	// Create a valid CryptoKey object
 	cryptoKeyMeta := &keys.CryptoKeyMeta{