Regen mocks and fix tests

Signed-off-by: Yuri Shkuro <github@ysh.us>

Author: yurishkuro

Makefile

@@ -418,6 +418,7 @@ install-mockery:
 .PHONY: generate-mocks
 generate-mocks: install-mockery
 	$(MOCKERY) -all -dir ./pkg/es/ -output ./pkg/es/mocks && rm pkg/es/mocks/ClientBuilder.go
+	$(MOCKERY) -all -dir ./storage/spanstore/ -output ./storage/spanstore/mocks
 
 .PHONY: echo-version
 echo-version:

cmd/ingester/app/processor/span_processor_test.go

@@ -15,6 +15,7 @@
 package processor
 
 import (
+	"context"
 	"errors"
 	"testing"
 
@@ -45,7 +46,7 @@ func TestSpanProcessor_Process(t *testing.T) {
 
 	message.On("Value").Return(data)
 	unmarshallerMock.On("Unmarshal", data).Return(span, nil)
-	writer.On("WriteSpan", span).Return(nil)
+	writer.On("WriteSpan", context.Background(), span).Return(nil)
 
 	assert.Nil(t, processor.Process(message))
 

cmd/opentelemetry/app/exporter/elasticsearchexporter/integration_test.go

@@ -38,6 +38,7 @@ import (
 	"github.com/jaegertracing/jaeger/plugin/storage/es"
 	"github.com/jaegertracing/jaeger/plugin/storage/es/spanstore/dbmodel"
 	"github.com/jaegertracing/jaeger/plugin/storage/integration"
+	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
 const (
@@ -61,13 +62,15 @@ type storageWrapper struct {
 	writer *esSpanWriter
 }
 
-func (s storageWrapper) WriteSpan(span *model.Span) error {
+var _ spanstore.Writer = (*storageWrapper)(nil)
+
+func (s storageWrapper) WriteSpan(ctx context.Context, span *model.Span) error {
 	// This fails because there is no binary tag type in OTEL and also OTEL span's status code is always created
 	//traces := jaegertranslator.ProtoBatchesToInternalTraces([]*model.Batch{{Process: span.Process, Spans: []*model.Span{span}}})
 	//_, err := s.writer.WriteTraces(context.Background(), traces)
 	converter := dbmodel.FromDomain{}
 	dbSpan := converter.FromDomainEmbedProcess(span)
-	_, err := s.writer.writeSpans(context.Background(), []*dbmodel.Span{dbSpan})
+	_, err := s.writer.writeSpans(ctx, []*dbmodel.Span{dbSpan})
 	return err
 }
 

cmd/query/app/grpc_handler_test.go

@@ -264,7 +264,7 @@ func TestArchiveTraceSuccessGRPC(t *testing.T) {
 	withServerAndClient(t, func(server *grpcServer, client *grpcClient) {
 		server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 			Return(mockTrace, nil).Once()
-		server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*model.Span")).
+		server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")).
 			Return(nil).Times(2)
 
 		_, err := client.ArchiveTrace(context.Background(), &api_v2.ArchiveTraceRequest{
@@ -295,7 +295,7 @@ func TestArchiveTraceFailureGRPC(t *testing.T) {
 
 		server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 			Return(mockTrace, nil).Once()
-		server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*model.Span")).
+		server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")).
 			Return(errStorageGRPC).Times(2)
 
 		_, err := client.ArchiveTrace(context.Background(), &api_v2.ArchiveTraceRequest{

cmd/query/app/handler_archive_test.go

@@ -103,7 +103,7 @@ func TestArchiveTrace_NoStorage(t *testing.T) {
 
 func TestArchiveTrace_Success(t *testing.T) {
 	mockWriter := &spanstoremocks.Writer{}
-	mockWriter.On("WriteSpan", mock.AnythingOfType("*model.Span")).
+	mockWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")).
 		Return(nil).Times(2)
 	withTestServer(t, func(ts *testServer) {
 		ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
@@ -116,7 +116,7 @@ func TestArchiveTrace_Success(t *testing.T) {
 
 func TestArchiveTrace_WriteErrors(t *testing.T) {
 	mockWriter := &spanstoremocks.Writer{}
-	mockWriter.On("WriteSpan", mock.AnythingOfType("*model.Span")).
+	mockWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")).
 		Return(errors.New("cannot save")).Times(2)
 	withTestServer(t, func(ts *testServer) {
 		ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).

cmd/query/app/querysvc/query_service_test.go

@@ -212,7 +212,7 @@ func TestArchiveTraceWithArchiveWriterError(t *testing.T) {
 	qs, readMock, _, _, writeMock := initializeTestServiceWithArchiveOptions()
 	readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(mockTrace, nil).Once()
-	writeMock.On("WriteSpan", mock.AnythingOfType("*model.Span")).
+	writeMock.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")).
 		Return(errors.New("cannot save")).Times(2)
 
 	type contextKey string
@@ -228,7 +228,7 @@ func TestArchiveTraceSuccess(t *testing.T) {
 	qs, readMock, _, _, writeMock := initializeTestServiceWithArchiveOptions()
 	readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(mockTrace, nil).Once()
-	writeMock.On("WriteSpan", mock.AnythingOfType("*model.Span")).
+	writeMock.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")).
 		Return(nil).Times(2)
 
 	type contextKey string

plugin/storage/grpc/shared/grpc_server_test.go

@@ -177,7 +177,7 @@ func TestGRPCServerFindTraceIDs(t *testing.T) {
 
 func TestGRPCServerWriteSpan(t *testing.T) {
 	withGRPCServer(func(r *grpcServerTest) {
-		r.impl.spanWriter.On("WriteSpan", &mockTraceSpans[0]).
+		r.impl.spanWriter.On("WriteSpan", context.Background(), &mockTraceSpans[0]).
 			Return(nil)
 
 		s, err := r.server.WriteSpan(context.Background(), &storage_v1.WriteSpanRequest{