fix

Signed-off-by: Yuri Shkuro <[email protected]>
jaegertracing · Jul 13, 2024 · 6b35c99 · 6b35c99
1 parent 487a09a
commit 6b35c99
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 111 deletions.
diff --git a/cmd/query/app/handler_archive_test.go b/cmd/query/app/handler_archive_test.go
@@ -48,7 +48,7 @@ func TestGetArchivedTrace_NotFound(t *testing.T) {
 	} {
 		tc := tc // capture loop var
 		t.Run(tc.name, func(t *testing.T) {
-			withTestServer(func(ts *testServer) {
+			withTestServer(t, func(ts *testServer) {
 				ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 					Return(nil, spanstore.ErrTraceNotFound).Once()
 				var response structuredResponse
@@ -66,7 +66,7 @@ func TestGetArchivedTraceSuccess(t *testing.T) {
 	mockReader := &spanstoremocks.Reader{}
 	mockReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(mockTrace, nil).Once()
-	withTestServer(func(ts *testServer) {
+	withTestServer(t, func(ts *testServer) {
 		// make main reader return NotFound
 		ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 			Return(nil, spanstore.ErrTraceNotFound).Once()
@@ -81,7 +81,7 @@ func TestGetArchivedTraceSuccess(t *testing.T) {
 
 // Test failure in parsing trace ID.
 func TestArchiveTrace_BadTraceID(t *testing.T) {
-	withTestServer(func(ts *testServer) {
+	withTestServer(t, func(ts *testServer) {
 		var response structuredResponse
 		err := postJSON(ts.server.URL+"/api/archive/badtraceid", []string{}, &response)
 		require.Error(t, err)
@@ -95,7 +95,7 @@ func TestArchiveTrace_TraceNotFound(t *testing.T) {
 		Return(nil, spanstore.ErrTraceNotFound).Once()
 	mockWriter := &spanstoremocks.Writer{}
 	// Not actually going to write the trace, so no need to define mockWriter action
-	withTestServer(func(ts *testServer) {
+	withTestServer(t, func(ts *testServer) {
 		// make main reader return NotFound
 		ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 			Return(nil, spanstore.ErrTraceNotFound).Once()
@@ -106,7 +106,7 @@ func TestArchiveTrace_TraceNotFound(t *testing.T) {
 }
 
 func TestArchiveTrace_NoStorage(t *testing.T) {
-	withTestServer(func(ts *testServer) {
+	withTestServer(t, func(ts *testServer) {
 		var response structuredResponse
 		err := postJSON(ts.server.URL+"/api/archive/"+mockTraceID.String(), []string{}, &response)
 		require.EqualError(t, err, `500 error from server: {"data":null,"total":0,"limit":0,"offset":0,"errors":[{"code":500,"msg":"archive span storage was not configured"}]}`+"\n")
@@ -117,7 +117,7 @@ func TestArchiveTrace_Success(t *testing.T) {
 	mockWriter := &spanstoremocks.Writer{}
 	mockWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")).
 		Return(nil).Times(2)
-	withTestServer(func(ts *testServer) {
+	withTestServer(t, func(ts *testServer) {
 		ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 			Return(mockTrace, nil).Once()
 		var response structuredResponse
@@ -130,7 +130,7 @@ func TestArchiveTrace_WriteErrors(t *testing.T) {
 	mockWriter := &spanstoremocks.Writer{}
 	mockWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")).
 		Return(errors.New("cannot save")).Times(2)
-	withTestServer(func(ts *testServer) {
+	withTestServer(t, func(ts *testServer) {
 		ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 			Return(mockTrace, nil).Once()
 		var response structuredResponse

diff --git a/cmd/query/app/handler_deps_test.go b/cmd/query/app/handler_deps_test.go
@@ -310,8 +310,7 @@ func TestFilterDependencies(t *testing.T) {
 }
 
 func TestGetDependenciesSuccess(t *testing.T) {
-	ts := initializeTestServer()
-	defer ts.server.Close()
+	ts := initializeTestServer(t)
 	expectedDependencies := []model.DependencyLink{{Parent: "killer", Child: "queen", CallCount: 12}}
 	endTs := time.Unix(0, 1476374248550*millisToNanosMultiplier)
 	ts.dependencyReader.On("GetDependencies",
@@ -332,29 +331,27 @@ func TestGetDependenciesSuccess(t *testing.T) {
 }
 
 func TestGetDependenciesCassandraFailure(t *testing.T) {
-	ts := initializeTestServer()
-	defer ts.server.Close()
+	ts := initializeTestServer(t)
 	endTs := time.Unix(0, 1476374248550*millisToNanosMultiplier)
-	ts.dependencyReader.On("GetDependencies", endTs, defaultDependencyLookbackDuration).Return(nil, errStorage).Times(1)
+	ts.dependencyReader.On("GetDependencies",
+		mock.Anything, // context
+		endTs,
+		defaultDependencyLookbackDuration).Return(nil, errStorage).Times(1)
 
 	var response structuredResponse
 	err := getJSON(ts.server.URL+"/api/dependencies?endTs=1476374248550&service=testing", &response)
 	require.Error(t, err)
 }
 
 func TestGetDependenciesEndTimeParsingFailure(t *testing.T) {
-	ts := initializeTestServer()
-	defer ts.server.Close()
-
+	ts := initializeTestServer(t)
 	var response structuredResponse
 	err := getJSON(ts.server.URL+"/api/dependencies?endTs=shazbot&service=testing", &response)
 	require.Error(t, err)
 }
 
 func TestGetDependenciesLookbackParsingFailure(t *testing.T) {
-	ts := initializeTestServer()
-	defer ts.server.Close()
-
+	ts := initializeTestServer(t)
 	var response structuredResponse
 	err := getJSON(ts.server.URL+"/api/dependencies?endTs=1476374248550&service=testing&lookback=shazbot", &response)
 	require.Error(t, err)