cleanup the generated test data for the schema

menutraversal_test/menu_traversal_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"flag"
+	"fmt"
 	"log"
 	"math/rand"
 	"os"
@@ -11,9 +12,11 @@ import (
 	"regexp"
 	"testing"
 
+	"git.grassecon.net/urdt/ussd/internal/storage"
 	"git.grassecon.net/urdt/ussd/internal/testutil"
 	"git.grassecon.net/urdt/ussd/internal/testutil/driver"
 	"github.com/gofrs/uuid"
+	"github.com/jackc/pgx/v5/pgxpool"
 )
 
 var (
@@ -25,6 +28,7 @@ var (
 
 var groupTestFile = flag.String("test-file", "group_test.json", "The test file to use for running the group tests")
 var database = flag.String("db", "gdbm", "Specify the database (gdbm or postgres)")
+var dbSchema = flag.String("schema", "test", "Specify the database schema (default test)")
 
 func testStore() string {
 	v, _ :=  filepath.Abs(".test_state/state.gdbm")
@@ -96,7 +100,29 @@ func TestMain(m *testing.M) {
 	}()
 
 	// Set the selected database
-	testutil.SetDatabase(*database)
+	testutil.SetDatabase(*database, *dbSchema)
+
+	// Cleanup the schema table after tests
+	defer func() {
+		if *database == "postgres" {
+			ctx := context.Background()
+			connStr := storage.BuildConnStr()
+			dbConn, err := pgxpool.New(ctx, connStr)
+			if err != nil {
+				log.Fatalf("Failed to connect to database for cleanup: %v", err)
+			}
+			defer dbConn.Close()
+
+			query := fmt.Sprintf("DELETE FROM %s.kv_vise;", *dbSchema)
+			_, execErr := dbConn.Exec(ctx, query)
+			if execErr != nil {
+				log.Printf("Failed to cleanup table %s.kv_vise: %v", *dbSchema, execErr)
+			} else {
+				log.Printf("Successfully cleaned up table %s.kv_vise", *dbSchema)
+			}
+		}
+	}()
+
 	m.Run()
 }