From ca1f69f04279a8ed04fbea91c05c41f30e8179a6 Mon Sep 17 00:00:00 2001
From: Antoine Grondin <antoinegrondin@gmail.com>
Date: Mon, 28 Oct 2024 18:41:27 +0900
Subject: [PATCH] fix(localstorage): test suite needs to allow closing storage

---
 internal/memstorage/memory.go  | 2 ++
 pkg/localstorage/queryable.go  | 1 +
 pkg/localstorage/test_suite.go | 7 ++++---
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/internal/memstorage/memory.go b/internal/memstorage/memory.go
index fc46c63..b1cd469 100644
--- a/internal/memstorage/memory.go
+++ b/internal/memstorage/memory.go
@@ -53,6 +53,8 @@ type SummarizedEvents struct {
 	}
 }
 
+func (str *MemStorage) Close() error { return nil }
+
 func (str *MemStorage) Query(ctx context.Context, q *typesv1.LogQuery) (<-chan localstorage.Cursor, error) {
 	if q.To != nil && q.From.AsTime().After(q.To.AsTime()) {
 		return nil, fmt.Errorf("invalid query, `to` is before `from`")
diff --git a/pkg/localstorage/queryable.go b/pkg/localstorage/queryable.go
index e4f44cd..f1b16f7 100644
--- a/pkg/localstorage/queryable.go
+++ b/pkg/localstorage/queryable.go
@@ -34,6 +34,7 @@ type Storage interface {
 	Queryable
 	SinkFor(ctx context.Context, machineID, sessionID int64) (_ sink.Sink, heartbeatIn time.Duration, _ error)
 	Heartbeat(ctx context.Context, machineID, sessionID int64) (time.Duration, error)
+	Close() error
 }
 
 type Queryable interface {
diff --git a/pkg/localstorage/test_suite.go b/pkg/localstorage/test_suite.go
index dacf1cd..888907b 100644
--- a/pkg/localstorage/test_suite.go
+++ b/pkg/localstorage/test_suite.go
@@ -181,10 +181,11 @@ func RunTest(t *testing.T, constructor func(t *testing.T) Storage) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 
-			mem := constructor(t)
+			db := constructor(t)
+			defer db.Close()
 
 			for _, leg := range tt.input {
-				snk, _, err := mem.SinkFor(ctx, leg.MachineId, leg.SessionId)
+				snk, _, err := db.SinkFor(ctx, leg.MachineId, leg.SessionId)
 				require.NoError(t, err)
 				for _, ev := range leg.Logs {
 					err = snk.Receive(ctx, ev)
@@ -201,7 +202,7 @@ func RunTest(t *testing.T, constructor func(t *testing.T) Storage) {
 				defer cancel()
 			}
 			now := time.Now()
-			cursors, err := mem.Query(queryctx, tt.q)
+			cursors, err := db.Query(queryctx, tt.q)
 			require.NoError(t, err)
 			got := drainCursors(t, queryctx, cursors)