From b0a773c3f8f8c3d9486fedf904ccfa08faf96d7a Mon Sep 17 00:00:00 2001
From: diamondburned <diamond@arikawa-hi.me>
Date: Sat, 4 Nov 2023 02:27:37 -0700
Subject: [PATCH] *: Improve WSDebug logging in tests

---
 gateway/gateway_test.go | 27 ++++++++++++++-------------
 voice/session_test.go   | 23 +++++++++++++----------
 2 files changed, 27 insertions(+), 23 deletions(-)

diff --git a/gateway/gateway_test.go b/gateway/gateway_test.go
index c3a6c429..4c902207 100644
--- a/gateway/gateway_test.go
+++ b/gateway/gateway_test.go
@@ -5,7 +5,6 @@ import (
 	"log"
 	"strconv"
 	"strings"
-	"sync"
 	"testing"
 	"time"
 
@@ -14,16 +13,18 @@ import (
 	"github.com/diamondburned/arikawa/v3/utils/ws"
 )
 
-var doLogOnce sync.Once
+func doLog(t *testing.T) {
+	if !testing.Verbose() {
+		return
+	}
 
-func doLog() {
-	doLogOnce.Do(func() {
-		if testing.Verbose() {
-			ws.WSDebug = func(v ...interface{}) {
-				log.Println(append([]interface{}{"Debug:"}, v...)...)
-			}
-		}
-	})
+	prev := ws.WSDebug
+	t.Cleanup(func() { ws.WSDebug = prev })
+
+	ws.WSDebug = func(v ...interface{}) {
+		t.Helper()
+		t.Log(v...)
+	}
 }
 
 func TestURL(t *testing.T) {
@@ -45,7 +46,7 @@ func TestURL(t *testing.T) {
 }
 
 func TestInvalidToken(t *testing.T) {
-	doLog()
+	doLog(t)
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
 	t.Cleanup(cancel)
@@ -90,7 +91,7 @@ func TestInvalidToken(t *testing.T) {
 }
 
 func TestIntegration(t *testing.T) {
-	doLog()
+	doLog(t)
 
 	config := testenv.Must(t)
 
@@ -108,7 +109,7 @@ func TestIntegration(t *testing.T) {
 }
 
 func TestReuseGateway(t *testing.T) {
-	doLog()
+	doLog(t)
 
 	config := testenv.Must(t)
 
diff --git a/voice/session_test.go b/voice/session_test.go
index b07c0554..4e338645 100644
--- a/voice/session_test.go
+++ b/voice/session_test.go
@@ -6,9 +6,6 @@ import (
 	"fmt"
 	"log"
 	"math/rand"
-	"os"
-	"runtime"
-	"strconv"
 	"sync"
 	"testing"
 	"time"
@@ -24,15 +21,18 @@ import (
 	"github.com/diamondburned/arikawa/v3/voice/voicegateway"
 )
 
-func TestMain(m *testing.M) {
-	ws.WSDebug = func(v ...interface{}) {
-		_, file, line, _ := runtime.Caller(1)
-		caller := file + ":" + strconv.Itoa(line)
-		log.Println(append([]interface{}{caller}, v...)...)
+func doLog(t *testing.T) {
+	if !testing.Verbose() {
+		return
 	}
 
-	code := m.Run()
-	os.Exit(code)
+	prev := ws.WSDebug
+	t.Cleanup(func() { ws.WSDebug = prev })
+
+	ws.WSDebug = func(v ...interface{}) {
+		t.Helper()
+		t.Log(v...)
+	}
 }
 
 type testState struct {
@@ -41,6 +41,9 @@ type testState struct {
 }
 
 func testOpen(t *testing.T) *testState {
+	t.Helper()
+	doLog(t)
+
 	config := testenv.Must(t)
 
 	s := state.New("Bot " + config.BotToken)