From af1f34861ea10b94db651b7e07f928285dd05934 Mon Sep 17 00:00:00 2001
From: Kegan Dougal <7190048+kegsay@users.noreply.github.com>
Date: Fri, 17 May 2024 16:09:02 +0100
Subject: [PATCH] Ensure txns are closed so we can wipe the db for other tests

---
 state/migrations/20231108122539_clear_stuck_invites_test.go | 1 +
 testutils/db.go                                             | 6 +++++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/state/migrations/20231108122539_clear_stuck_invites_test.go b/state/migrations/20231108122539_clear_stuck_invites_test.go
index 616fd62..a4b4257 100644
--- a/state/migrations/20231108122539_clear_stuck_invites_test.go
+++ b/state/migrations/20231108122539_clear_stuck_invites_test.go
@@ -151,6 +151,7 @@ func TestClearStuckInvites(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer tx.Rollback()
 
 	// users in room B (bob) and F (doris) should be reset.
 	tokens, err := tokensTable.TokenForEachDevice(tx)
diff --git a/testutils/db.go b/testutils/db.go
index d4106aa..af79102 100644
--- a/testutils/db.go
+++ b/testutils/db.go
@@ -1,11 +1,13 @@
 package testutils
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"os"
 	"os/exec"
 	"os/user"
+	"time"
 )
 
 var Quiet = false
@@ -64,7 +66,9 @@ func PrepareDBConnectionString() (connStr string) {
 	if err != nil {
 		panic(err)
 	}
-	_, err = db.Exec(`DROP SCHEMA public CASCADE;
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+	_, err = db.ExecContext(ctx, `DROP SCHEMA public CASCADE;
 		CREATE SCHEMA public;`)
 	if err != nil {
 		panic(err)