From 98fa9137f76d51e6d8c64550b640a17ef706f2fa Mon Sep 17 00:00:00 2001
From: Mahad Zaryab <mahadzaryab1@gmail.com>
Date: Sun, 23 Feb 2025 10:20:06 -0500
Subject: [PATCH] Fix Tests

Signed-off-by: Mahad Zaryab <mahadzaryab1@gmail.com>
---
 .../storage/v2/v1adapter/spanreader_test.go   | 18 ++++----
 .../storage/v2/v1adapter/tracereader_test.go  | 42 ++++++++++---------
 2 files changed, 34 insertions(+), 26 deletions(-)

diff --git a/internal/storage/v2/v1adapter/spanreader_test.go b/internal/storage/v2/v1adapter/spanreader_test.go
index ec1052c2bcc..92950b53167 100644
--- a/internal/storage/v2/v1adapter/spanreader_test.go
+++ b/internal/storage/v2/v1adapter/spanreader_test.go
@@ -336,7 +336,7 @@ func TestSpanReader_FindTraceIDs(t *testing.T) {
 		name             string
 		query            *spanstore.TraceQueryParameters
 		expectedQuery    tracestore.TraceQueryParams
-		traceIDs         []pcommon.TraceID
+		getTraceParams   []tracestore.GetTraceParams
 		expectedTraceIDs []model.TraceID
 		err              error
 		expectedErr      error
@@ -360,7 +360,7 @@ func TestSpanReader_FindTraceIDs(t *testing.T) {
 			expectedQuery: tracestore.TraceQueryParams{
 				ServiceName: "service1",
 			},
-			traceIDs:         []pcommon.TraceID{},
+			getTraceParams:   []tracestore.GetTraceParams{},
 			expectedTraceIDs: nil,
 		},
 		{
@@ -371,9 +371,13 @@ func TestSpanReader_FindTraceIDs(t *testing.T) {
 			expectedQuery: tracestore.TraceQueryParams{
 				ServiceName: "service1",
 			},
-			traceIDs: []pcommon.TraceID{
-				pcommon.TraceID([16]byte{0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2}),
-				pcommon.TraceID([16]byte{0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4}),
+			getTraceParams: []tracestore.GetTraceParams{
+				{
+					TraceID: [16]byte{0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2},
+				},
+				{
+					TraceID: [16]byte{0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4},
+				},
 			},
 			expectedTraceIDs: []model.TraceID{
 				model.NewTraceID(1, 2),
@@ -385,8 +389,8 @@ func TestSpanReader_FindTraceIDs(t *testing.T) {
 	for _, test := range tests {
 		tr := tracestoremocks.Reader{}
 		tr.On("FindTraceIDs", mock.Anything, test.expectedQuery).
-			Return(iter.Seq2[[]pcommon.TraceID, error](func(yield func([]pcommon.TraceID, error) bool) {
-				yield(test.traceIDs, test.err)
+			Return(iter.Seq2[[]tracestore.GetTraceParams, error](func(yield func([]tracestore.GetTraceParams, error) bool) {
+				yield(test.getTraceParams, test.err)
 			})).Once()
 
 		sr := NewSpanReader(&tr)
diff --git a/internal/storage/v2/v1adapter/tracereader_test.go b/internal/storage/v2/v1adapter/tracereader_test.go
index 67af5c6a05c..59b2d98dbd2 100644
--- a/internal/storage/v2/v1adapter/tracereader_test.go
+++ b/internal/storage/v2/v1adapter/tracereader_test.go
@@ -402,10 +402,10 @@ func TestTraceReader_FindTracesEarlyStop(t *testing.T) {
 
 func TestTraceReader_FindTraceIDsDelegatesResponse(t *testing.T) {
 	tests := []struct {
-		name             string
-		modelTraceIDs    []model.TraceID
-		expectedTraceIDs []pcommon.TraceID
-		err              error
+		name           string
+		modelTraceIDs  []model.TraceID
+		expectedParams []tracestore.GetTraceParams
+		err            error
 	}{
 		{
 			name: "successful response",
@@ -413,26 +413,30 @@ func TestTraceReader_FindTraceIDsDelegatesResponse(t *testing.T) {
 				{Low: 3, High: 2},
 				{Low: 4, High: 3},
 			},
-			expectedTraceIDs: []pcommon.TraceID{
-				pcommon.TraceID([]byte{0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3}),
-				pcommon.TraceID([]byte{0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4}),
+			expectedParams: []tracestore.GetTraceParams{
+				{
+					TraceID: pcommon.TraceID([]byte{0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3}),
+				},
+				{
+					TraceID: pcommon.TraceID([]byte{0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4}),
+				},
 			},
 		},
 		{
-			name:             "empty response",
-			modelTraceIDs:    []model.TraceID{},
-			expectedTraceIDs: nil,
+			name:           "empty response",
+			modelTraceIDs:  []model.TraceID{},
+			expectedParams: nil,
 		},
 		{
-			name:             "nil response",
-			modelTraceIDs:    nil,
-			expectedTraceIDs: nil,
+			name:           "nil response",
+			modelTraceIDs:  nil,
+			expectedParams: nil,
 		},
 		{
-			name:             "error response",
-			modelTraceIDs:    nil,
-			expectedTraceIDs: nil,
-			err:              errors.New("test error"),
+			name:           "error response",
+			modelTraceIDs:  nil,
+			expectedParams: nil,
+			err:            errors.New("test error"),
 		},
 	}
 	for _, test := range tests {
@@ -456,7 +460,7 @@ func TestTraceReader_FindTraceIDsDelegatesResponse(t *testing.T) {
 			traceReader := &TraceReader{
 				spanReader: sr,
 			}
-			traceIDs, err := jiter.FlattenWithErrors(traceReader.FindTraceIDs(
+			getTraceParams, err := jiter.FlattenWithErrors(traceReader.FindTraceIDs(
 				context.Background(),
 				tracestore.TraceQueryParams{
 					ServiceName:   "service",
@@ -470,7 +474,7 @@ func TestTraceReader_FindTraceIDsDelegatesResponse(t *testing.T) {
 				},
 			))
 			require.ErrorIs(t, err, test.err)
-			require.Equal(t, test.expectedTraceIDs, traceIDs)
+			require.Equal(t, test.expectedParams, getTraceParams)
 		})
 	}
 }