Improve test using two mocks

Signed-off-by: Pavol Loffay <ploffay@redhat.com>

plugin/storage/es/spanstore/reader_test.go

@@ -205,11 +205,14 @@ func TestSpanReader_GetTrace(t *testing.T) {
 	})
 }
 
-func TestSpanReader_multiRead_followUp_traceId(t *testing.T) {
+func TestSpanReader_multiRead_followUp_query(t *testing.T) {
 	withSpanReader(func(r *spanReaderTest) {
 		now := time.Now()
-		span := dbmodel.Span{SpanID: "0", TraceID:"1", StartTime: model.TimeAsEpochMicroseconds(now)}
-		spanBytes, err := json.Marshal(span)
+		spanID1 := dbmodel.Span{SpanID: "0", TraceID:"1", StartTime: model.TimeAsEpochMicroseconds(now)}
+		spanBytesID1, err := json.Marshal(spanID1)
+		require.NoError(t, err)
+		spanID2 := dbmodel.Span{SpanID: "0", TraceID:"2", StartTime: model.TimeAsEpochMicroseconds(now)}
+		spanBytesID2, err := json.Marshal(spanID2)
 		require.NoError(t, err)
 
 		id1Query := elastic.NewTermQuery("traceID", model.TraceID{High: 0, Low:1}.String())
@@ -225,39 +228,56 @@ func TestSpanReader_multiRead_followUp_traceId(t *testing.T) {
 		id1SearchSpanTime := elastic.NewSearchRequest().
 			IgnoreUnavailable(true).
 			Type(spanType).
-			Source(r.reader.sourceFn(id1Query, span.StartTime))
+			Source(r.reader.sourceFn(id1Query, spanID1.StartTime))
 
 		multiSearchService := &mocks.MultiSearchService{}
-		multiSearchService.On("Add", id1Search, id2Search).Return(multiSearchService)
-		multiSearchService.On("Add", id1SearchSpanTime).Return(multiSearchService)
-		multiSearchService.On("Index", mock.AnythingOfType("string"), mock.AnythingOfType("string"),
-			mock.AnythingOfType("string"), mock.AnythingOfType("string"), mock.AnythingOfType("string")).Return(multiSearchService)
+		firstMultiSearch := &mocks.MultiSearchService{}
+		secondMultiSearch := &mocks.MultiSearchService{}
+		multiSearchService.On("Add", id1Search, id2Search).Return(firstMultiSearch)
+		multiSearchService.On("Add", id1SearchSpanTime).Return(secondMultiSearch)
+
+		firstMultiSearch.On("Index", mock.AnythingOfType("string")).Return(firstMultiSearch)
+		secondMultiSearch.On("Index", mock.AnythingOfType("string")).Return(secondMultiSearch)
 		r.client.On("MultiSearch").Return(multiSearchService)
-		multiSearchMock := multiSearchService.On("Do", mock.AnythingOfType("*context.emptyCtx"))
 
-		hits := []*elastic.SearchHit{
-			{Source: (*json.RawMessage)(&spanBytes)},
-		}
+		fistMultiSearchMock := firstMultiSearch.On("Do", mock.AnythingOfType("*context.emptyCtx"))
+		secondMultiSearchMock := secondMultiSearch.On("Do", mock.AnythingOfType("*context.emptyCtx"))
+
 		// set TotalHits to two to trigger the follow up query
 		// the client will return only one span therefore the implementation
 		// triggers follow up query for the same traceID with the timestamp of the last span
-		searchHits := &elastic.SearchHits{Hits: hits, TotalHits: 2}
-		multiSearchMock.
+		searchHitsID1 := &elastic.SearchHits{Hits: []*elastic.SearchHit{
+			{Source: (*json.RawMessage)(&spanBytesID1)},
+		}, TotalHits: 2}
+		fistMultiSearchMock.
 			Return(&elastic.MultiSearchResult{
 				Responses: []*elastic.SearchResult{
-					{Hits: searchHits},
+					{Hits: searchHitsID1},
+				},
+			}, nil)
+
+		searchHitsID2 := &elastic.SearchHits{Hits: []*elastic.SearchHit{
+			{Source: (*json.RawMessage)(&spanBytesID2)},
+		}, TotalHits: 1}
+		secondMultiSearchMock.
+			Return(&elastic.MultiSearchResult{
+				Responses: []*elastic.SearchResult{
+					{Hits: searchHitsID2},
 				},
 			}, nil)
 
 		traces, err := r.reader.multiRead(context.Background(), []model.TraceID{{High:0, Low:1}, {High:0, Low:2}}, now, now)
 		require.NoError(t, err)
 		require.NotNil(t, traces)
+		require.Len(t, traces, 2)
 
-		expectedSpans, err := r.reader.collectSpans(hits)
+		toDomain := dbmodel.NewToDomain("-")
+		sModel1, err := toDomain.SpanToDomain(&spanID1)
 		require.NoError(t, err)
-
-		require.Len(t, traces, 1)
-		assert.EqualValues(t, traces[0].Spans[0], expectedSpans[0])
+		sModel2, err := toDomain.SpanToDomain(&spanID2)
+		require.NoError(t, err)
+		assert.EqualValues(t, traces[0].Spans[0], sModel1)
+		assert.EqualValues(t, traces[1].Spans[0], sModel2)
 	})
 }