From 57bcd5aa7fe777e515624dddef8082d398a91c00 Mon Sep 17 00:00:00 2001
From: Jimmy Moore <jamesmoore@loopholelabs.io>
Date: Fri, 1 Nov 2024 12:51:28 +0000
Subject: [PATCH] Updated with threadsafe bytes buffer for log in tests

Signed-off-by: Jimmy Moore <jamesmoore@loopholelabs.io>
---
 pkg/storage/device/device_sync_test.go |  3 +--
 pkg/storage/expose/nbd_test.go         |  4 ++--
 pkg/testutils/logging_output.go        | 29 ++++++++++++++++++++++++++
 3 files changed, 32 insertions(+), 4 deletions(-)
 create mode 100644 pkg/testutils/logging_output.go

diff --git a/pkg/storage/device/device_sync_test.go b/pkg/storage/device/device_sync_test.go
index db1c0b3a..90f6d065 100644
--- a/pkg/storage/device/device_sync_test.go
+++ b/pkg/storage/device/device_sync_test.go
@@ -1,7 +1,6 @@
 package device
 
 import (
-	"bytes"
 	"crypto/rand"
 	"crypto/sha256"
 	"fmt"
@@ -52,7 +51,7 @@ func TestDeviceSync(t *testing.T) {
 	err := s.Decode([]byte(testSyncSchema))
 	assert.NoError(t, err)
 
-	logBuffer := &bytes.Buffer{}
+	logBuffer := &testutils.SafeWriteBuffer{}
 	l := logging.New(logging.Zerolog, "device", logBuffer)
 	l.SetLevel(types.TraceLevel)
 
diff --git a/pkg/storage/expose/nbd_test.go b/pkg/storage/expose/nbd_test.go
index 45920bbf..ca71592e 100644
--- a/pkg/storage/expose/nbd_test.go
+++ b/pkg/storage/expose/nbd_test.go
@@ -1,7 +1,6 @@
 package expose
 
 import (
-	"bytes"
 	"crypto/rand"
 	"fmt"
 	"os"
@@ -14,6 +13,7 @@ import (
 	"github.com/loopholelabs/logging/types"
 	"github.com/loopholelabs/silo/pkg/storage/modules"
 	"github.com/loopholelabs/silo/pkg/storage/sources"
+	"github.com/loopholelabs/silo/pkg/testutils"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -36,7 +36,7 @@ func TestNBDNLDevice(t *testing.T) {
 	size := 4096 * 1024 * 1024
 	prov := sources.NewMemoryStorage(size)
 
-	logBuffer := &bytes.Buffer{}
+	logBuffer := &testutils.SafeWriteBuffer{}
 	// Enable logging here to make sure it doesn't break anything...
 	l := logging.New(logging.Zerolog, "silo", logBuffer)
 	l.SetLevel(types.TraceLevel)
diff --git a/pkg/testutils/logging_output.go b/pkg/testutils/logging_output.go
new file mode 100644
index 00000000..f18100ff
--- /dev/null
+++ b/pkg/testutils/logging_output.go
@@ -0,0 +1,29 @@
+package testutils
+
+import (
+	"bytes"
+	"sync"
+)
+
+type SafeWriteBuffer struct {
+	bufferLock sync.Mutex
+	buffer     bytes.Buffer
+}
+
+func (swb *SafeWriteBuffer) Write(p []byte) (n int, err error) {
+	swb.bufferLock.Lock()
+	defer swb.bufferLock.Unlock()
+	return swb.buffer.Write(p)
+}
+
+func (swb *SafeWriteBuffer) Bytes() []byte {
+	swb.bufferLock.Lock()
+	defer swb.bufferLock.Unlock()
+	return swb.buffer.Bytes()
+}
+
+func (swb *SafeWriteBuffer) Len() int {
+	swb.bufferLock.Lock()
+	defer swb.bufferLock.Unlock()
+	return swb.buffer.Len()
+}