add more tests

Signed-off-by: Zhang Xin <[email protected]>
jaegertracing · Dec 20, 2024 · 55ea627 · 55ea627
1 parent f188a58
commit 55ea627
Showing 1 changed file with 35 additions and 4 deletions.
diff --git a/cmd/query/app/handler_archive_test.go b/cmd/query/app/handler_archive_test.go
@@ -7,6 +7,7 @@ package app
 import (
 	"errors"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -67,6 +68,29 @@ func TestGetArchivedTraceSuccess(t *testing.T) {
 	}, querysvc.QueryServiceOptions{ArchiveSpanReader: mockReader})
 }
 
+func TestGetArchivedTraceWithTimeWindowSuccess(t *testing.T) {
+	traceID := model.NewTraceID(0, 123456)
+	mockReader := &spanstoremocks.Reader{}
+	expectedQuery := spanstore.GetTraceParameters{
+		TraceID:   traceID,
+		StartTime: time.UnixMicro(1),
+		EndTime:   time.UnixMicro(2),
+	}
+	mockReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedQuery).
+		Return(mockTrace, nil).Once()
+	withTestServer(t, func(ts *testServer) {
+		// make main reader return NotFound
+		ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")).
+			Return(nil, spanstore.ErrTraceNotFound).Once()
+		var response structuredTraceResponse
+		err := getJSON(ts.server.URL+"/api/traces/"+mockTraceID.String()+"?start=1&end=2", &response)
+		require.NoError(t, err)
+		assert.Empty(t, response.Errors)
+		assert.Len(t, response.Traces, 1)
+		assert.Equal(t, traceID.String(), string(response.Traces[0].TraceID))
+	}, querysvc.QueryServiceOptions{ArchiveSpanReader: mockReader})
+}
+
 // Test failure in parsing trace ID.
 func TestArchiveTrace_BadTraceID(t *testing.T) {
 	withTestServer(t, func(ts *testServer) {
@@ -128,6 +152,9 @@ func TestArchiveTrace_WriteErrors(t *testing.T) {
 }
 
 func TestArchiveTrace_BadTimeWindow(t *testing.T) {
+	mockReader := &spanstoremocks.Reader{}
+	mockReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")).
+		Return(mockTrace, nil).Times(2)
 	testCases := []struct {
 		name  string
 		query string
@@ -143,10 +170,14 @@ func TestArchiveTrace_BadTimeWindow(t *testing.T) {
 	}
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			ts := initializeTestServer(t)
-			var response structuredResponse
-			err := getJSON(ts.server.URL+"/api/archive/"+mockTraceID.String()+"?"+tc.query, &response)
-			require.Error(t, err)
+			withTestServer(t, func(ts *testServer) {
+				// make main reader return NotFound
+				ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")).
+					Return(nil, spanstore.ErrTraceNotFound).Once()
+				var response structuredTraceResponse
+				err := getJSON(ts.server.URL+"/api/traces/"+mockTraceID.String()+"?"+tc.query, &response)
+				require.Error(t, err)
+			}, querysvc.QueryServiceOptions{ArchiveSpanReader: mockReader})
 		})
 	}
 }