From 45d2de231e7209110cdac472cc66f17afb0422eb Mon Sep 17 00:00:00 2001 From: Yuri Shkuro Date: Wed, 22 Mar 2023 21:10:15 -0400 Subject: [PATCH] [test] Use correct Context instance with mocks (#4332) ## Which problem is this PR solving? - Should fix #4327 ## Short description of the changes - Test on tip were failing due to mismatch of context types passed to mocks. --------- Signed-off-by: Yuri Shkuro --- cmd/ingester/app/processor/span_processor_test.go | 2 +- plugin/storage/es/spanstore/reader_test.go | 13 +++++-------- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/cmd/ingester/app/processor/span_processor_test.go b/cmd/ingester/app/processor/span_processor_test.go index 5df0159b7a6..2e61c16e09f 100644 --- a/cmd/ingester/app/processor/span_processor_test.go +++ b/cmd/ingester/app/processor/span_processor_test.go @@ -49,7 +49,7 @@ func TestSpanProcessor_Process(t *testing.T) { message.On("Value").Return(data) mockUnmarshaller.On("Unmarshal", data).Return(span, nil) - mockWriter.On("WriteSpan", context.Background(), span). + mockWriter.On("WriteSpan", context.TODO(), span). Return(nil). Run(func(args mock.Arguments) { span := args[1].(*model.Span) diff --git a/plugin/storage/es/spanstore/reader_test.go b/plugin/storage/es/spanstore/reader_test.go index f5366b7a3bd..c581a423d1e 100644 --- a/plugin/storage/es/spanstore/reader_test.go +++ b/plugin/storage/es/spanstore/reader_test.go @@ -356,8 +356,8 @@ func TestSpanReader_multiRead_followUp_query(t *testing.T) { secondMultiSearch.On("Index", mock.AnythingOfType("string")).Return(secondMultiSearch) r.client.On("MultiSearch").Return(multiSearchService) - fistMultiSearchMock := firstMultiSearch.On("Do", mock.AnythingOfType("*context.emptyCtx")) - secondMultiSearchMock := secondMultiSearch.On("Do", mock.AnythingOfType("*context.emptyCtx")) + fistMultiSearchMock := firstMultiSearch.On("Do", mock.Anything) + secondMultiSearchMock := secondMultiSearch.On("Do", mock.Anything) // set TotalHits to two to trigger the follow up query // the client will return only one span therefore the implementation @@ -951,7 +951,7 @@ func mockMultiSearchService(r *spanReaderTest) *mock.Call { multiSearchService.On("Index", mock.AnythingOfType("string"), mock.AnythingOfType("string"), mock.AnythingOfType("string"), mock.AnythingOfType("string"), mock.AnythingOfType("string")).Return(multiSearchService) r.client.On("MultiSearch").Return(multiSearchService) - return multiSearchService.On("Do", mock.AnythingOfType("*context.valueCtx")) + return multiSearchService.On("Do", mock.Anything) } func mockArchiveMultiSearchService(r *spanReaderTest, indexName string) *mock.Call { @@ -959,7 +959,7 @@ func mockArchiveMultiSearchService(r *spanReaderTest, indexName string) *mock.Ca multiSearchService.On("Add", mock.Anything, mock.Anything, mock.Anything).Return(multiSearchService) multiSearchService.On("Index", indexName).Return(multiSearchService) r.client.On("MultiSearch").Return(multiSearchService) - return multiSearchService.On("Do", mock.AnythingOfType("*context.valueCtx")) + return multiSearchService.On("Do", mock.Anything) } // matchTermsAggregation uses reflection to match the size attribute of the TermsAggregation; neither @@ -981,10 +981,7 @@ func mockSearchService(r *spanReaderTest) *mock.Call { searchService.On("Aggregation", stringMatcher(operationsAggregation), mock.MatchedBy(matchTermsAggregation)).Return(searchService) searchService.On("Aggregation", stringMatcher(traceIDAggregation), mock.AnythingOfType("*elastic.TermsAggregation")).Return(searchService) r.client.On("Search", mock.AnythingOfType("string"), mock.AnythingOfType("string"), mock.AnythingOfType("string"), mock.AnythingOfType("string"), mock.AnythingOfType("string")).Return(searchService) - return searchService.On("Do", mock.MatchedBy(func(ctx context.Context) bool { - t := reflect.TypeOf(ctx).String() - return t == "*context.valueCtx" || t == "*context.emptyCtx" - })) + return searchService.On("Do", mock.Anything) } func TestTraceQueryParameterValidation(t *testing.T) {