Get to 100pc coverage

Signed-off-by: Prithvi Raj <[email protected]>

Author: vprithvi

Diff for: cmd/query/app/handler_test.go

@@ -133,7 +133,7 @@ func withTestServer(t *testing.T, doTest func(s *testServer), options ...Handler
 func TestGetTraceSuccess(t *testing.T) {
 	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	readMock.On("GetTrace", mock.Anything, mock.AnythingOfType("model.TraceID")).
+	readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(mockTrace, nil).Once()
 
 	var response structuredResponse
@@ -273,7 +273,7 @@ func TestGetTrace(t *testing.T) {
 func TestGetTraceDBFailure(t *testing.T) {
 	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	readMock.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+	readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(nil, errStorage).Once()
 
 	var response structuredResponse
@@ -462,9 +462,9 @@ func TestGetServicesSuccess(t *testing.T) {
 }
 
 func TestGetServicesStorageFailure(t *testing.T) {
-	server, mock, _ := initializeTestServer()
+	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	mock.On("GetServices").Return(nil, errStorage).Once()
+	readMock.On("GetServices", mock.AnythingOfType("*context.valueCtx")).Return(nil, errStorage).Once()
 
 	var response structuredResponse
 	err := getJSON(server.URL+"/api/services", &response)
@@ -497,9 +497,9 @@ func TestGetOperationsNoServiceName(t *testing.T) {
 }
 
 func TestGetOperationsStorageFailure(t *testing.T) {
-	server, mock, _ := initializeTestServer()
+	server, reaMock, _ := initializeTestServer()
 	defer server.Close()
-	mock.On("GetOperations", "trifle").Return(nil, errStorage).Once()
+	reaMock.On("GetOperations", mock.AnythingOfType("*context.valueCtx"), "trifle").Return(nil, errStorage).Once()
 
 	var response structuredResponse
 	err := getJSON(server.URL+"/api/operations?service=trifle", &response)
@@ -523,9 +523,9 @@ func TestGetOperationsLegacySuccess(t *testing.T) {
 }
 
 func TestGetOperationsLegacyStorageFailure(t *testing.T) {
-	server, mock, _ := initializeTestServer()
+	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	mock.On("GetOperations", "trifle").Return(nil, errStorage).Once()
+	readMock.On("GetOperations", mock.AnythingOfType("*context.valueCtx"), "trifle").Return(nil, errStorage).Once()
 
 	var response structuredResponse
 	err := getJSON(server.URL+"/api/services/trifle/operations", &response)

Diff for: storage/spanstore/metrics/decorator_test.go

@@ -41,6 +41,8 @@ func TestSuccessfulUnderlyingCalls(t *testing.T) {
 	mrs.GetTrace(context.Background(), model.TraceID{})
 	mockReader.On("FindTraces", context.Background(), &spanstore.TraceQueryParameters{}).Return([]*model.Trace{}, nil)
 	mrs.FindTraces(context.Background(), &spanstore.TraceQueryParameters{})
+	mockReader.On("FindTraceIDs", context.Background(), &spanstore.TraceQueryParameters{}).Return([]model.TraceID{}, nil)
+	mrs.FindTraceIDs(context.Background(), &spanstore.TraceQueryParameters{})
 	counters, gauges := mf.Snapshot()
 	expecteds := map[string]int64{
 		"requests|operation=get_operations|result=ok":  1,
@@ -49,6 +51,8 @@ func TestSuccessfulUnderlyingCalls(t *testing.T) {
 		"requests|operation=get_trace|result=err":      0,
 		"requests|operation=find_traces|result=ok":     1,
 		"requests|operation=find_traces|result=err":    0,
+		"requests|operation=find_trace_ids|result=ok":  1,
+		"requests|operation=find_trace_ids|result=err": 0,
 		"requests|operation=get_services|result=ok":    1,
 		"requests|operation=get_services|result=err":   0,
 	}
@@ -94,6 +98,8 @@ func TestFailingUnderlyingCalls(t *testing.T) {
 	mrs.GetTrace(context.Background(), model.TraceID{})
 	mockReader.On("FindTraces", context.Background(), &spanstore.TraceQueryParameters{}).Return(nil, errors.New("Failure"))
 	mrs.FindTraces(context.Background(), &spanstore.TraceQueryParameters{})
+	mockReader.On("FindTraceIDs", context.Background(), &spanstore.TraceQueryParameters{}).Return(nil, errors.New("Failure"))
+	mrs.FindTraceIDs(context.Background(), &spanstore.TraceQueryParameters{})
 	counters, gauges := mf.Snapshot()
 	expecteds := map[string]int64{
 		"requests|operation=get_operations|result=ok":  0,
@@ -102,6 +108,8 @@ func TestFailingUnderlyingCalls(t *testing.T) {
 		"requests|operation=get_trace|result=err":      1,
 		"requests|operation=find_traces|result=ok":     0,
 		"requests|operation=find_traces|result=err":    1,
+		"requests|operation=find_trace_ids|result=ok":   0,
+		"requests|operation=find_trace_ids|result=err":  1,
 		"requests|operation=get_services|result=ok":    0,
 		"requests|operation=get_services|result=err":   1,
 	}