From 2e237868aee872750e53a38f08c1fac03c88237c Mon Sep 17 00:00:00 2001
From: Alok Kumar Singh <dev.alok.singh123@gmail.com>
Date: Mon, 9 Dec 2024 21:50:31 +0530
Subject: [PATCH] Added sampling store mock factory in tests

Signed-off-by: Alok Kumar Singh <dev.alok.singh123@gmail.com>
---
 cmd/remote-storage/app/server.go               |  7 ++++++-
 cmd/remote-storage/app/server_test.go          | 17 +++++++++++++----
 cmd/remote-storage/main.go                     |  7 ++++++-
 plugin/storage/grpc/shared/grpc_client_test.go |  2 +-
 4 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/cmd/remote-storage/app/server.go b/cmd/remote-storage/app/server.go
index e4b9c081bf1..af19ec282e6 100644
--- a/cmd/remote-storage/app/server.go
+++ b/cmd/remote-storage/app/server.go
@@ -26,6 +26,10 @@ import (
 	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
+const (
+	DEFAULT_MAX_BUCKET_SIZE = 1
+)
+
 // Server runs a gRPC server
 type Server struct {
 	opts *Options
@@ -68,7 +72,8 @@ func createGRPCHandler(f storage.BaseFactory, samplingStoreFactory storage.Sampl
 	if err != nil {
 		return nil, err
 	}
-	samplingStore, err := samplingStoreFactory.CreateSamplingStore(1)
+	// TODO: Update this to use bucket size from config
+	samplingStore, err := samplingStoreFactory.CreateSamplingStore(DEFAULT_MAX_BUCKET_SIZE)
 	if err != nil {
 		return nil, err
 	}
diff --git a/cmd/remote-storage/app/server_test.go b/cmd/remote-storage/app/server_test.go
index cf52b3633b0..77cc400914e 100644
--- a/cmd/remote-storage/app/server_test.go
+++ b/cmd/remote-storage/app/server_test.go
@@ -41,13 +41,14 @@ func TestNewServer_CreateStorageErrors(t *testing.T) {
 	factory.On("CreateSpanWriter").Return(nil, nil)
 	factory.On("CreateDependencyReader").Return(nil, errors.New("no deps")).Once()
 	factory.On("CreateDependencyReader").Return(nil, nil)
+	samplingStoreFactoryMocks := factoryMocks.NewSamplingStoreFactory(t)
 	f := func() (*Server, error) {
 		return NewServer(
 			&Options{GRPCHostPort: ":0"},
 			factory,
 			tenancy.NewManager(&tenancy.Options{}),
 			telemetry.NoopSettings(),
-			nil,
+			samplingStoreFactoryMocks,
 		)
 	}
 	_, err := f()
@@ -114,12 +115,14 @@ func TestNewServer_TLSConfigError(t *testing.T) {
 		ReportStatus: telemetry.HCAdapter(healthcheck.New()),
 	}
 	storageMocks := newStorageMocks()
+	samplingStoreFactoryMocks := factoryMocks.NewSamplingStoreFactory(t)
+
 	_, err := NewServer(
 		&Options{GRPCHostPort: ":8081", TLSGRPC: tlsCfg},
 		storageMocks.factory,
 		tenancy.NewManager(&tenancy.Options{}),
 		telset,
-		nil,
+		samplingStoreFactoryMocks,
 	)
 	assert.ErrorContains(t, err, "invalid TLS config")
 }
@@ -322,12 +325,15 @@ func TestServerGRPCTLS(t *testing.T) {
 				Logger:       flagsSvc.Logger,
 				ReportStatus: telemetry.HCAdapter(flagsSvc.HC()),
 			}
+
+			samplingStoreFactoryMocks := factoryMocks.NewSamplingStoreFactory(t)
+
 			server, err := NewServer(
 				serverOptions,
 				storageMocks.factory,
 				tm,
 				telset,
-				nil,
+				samplingStoreFactoryMocks,
 			)
 			require.NoError(t, err)
 			require.NoError(t, server.Start())
@@ -372,12 +378,15 @@ func TestServerHandlesPortZero(t *testing.T) {
 		Logger:       flagsSvc.Logger,
 		ReportStatus: telemetry.HCAdapter(flagsSvc.HC()),
 	}
+
+	samplingStoreFactoryMocks := factoryMocks.NewSamplingStoreFactory(t)
+
 	server, err := NewServer(
 		&Options{GRPCHostPort: ":0"},
 		storageMocks.factory,
 		tenancy.NewManager(&tenancy.Options{}),
 		telset,
-		nil,
+		samplingStoreFactoryMocks,
 	)
 	require.NoError(t, err)
 
diff --git a/cmd/remote-storage/main.go b/cmd/remote-storage/main.go
index dacd5a7aa66..d52b7df7270 100644
--- a/cmd/remote-storage/main.go
+++ b/cmd/remote-storage/main.go
@@ -74,10 +74,15 @@ func main() {
 				logger.Fatal("Failed to init storage factory", zap.Error(err))
 			}
 
+			samplingStoreFactory, err := storageFactory.CreateSamplingStoreFactory()
+			if err != nil {
+				logger.Fatal("Failed to init sampling storage factory", zap.Error(err))
+			}
+
 			tm := tenancy.NewManager(&opts.Tenancy)
 			telset := baseTelset // copy
 			telset.Metrics = metricsFactory
-			server, err := app.NewServer(opts, storageFactory, tm, telset, nil)
+			server, err := app.NewServer(opts, storageFactory, tm, telset, samplingStoreFactory)
 			if err != nil {
 				logger.Fatal("Failed to create server", zap.Error(err))
 			}
diff --git a/plugin/storage/grpc/shared/grpc_client_test.go b/plugin/storage/grpc/shared/grpc_client_test.go
index a0e97bc2cad..0b044aaee21 100644
--- a/plugin/storage/grpc/shared/grpc_client_test.go
+++ b/plugin/storage/grpc/shared/grpc_client_test.go
@@ -112,7 +112,7 @@ func TestNewGRPCClient(t *testing.T) {
 	assert.Implements(t, (*storage_v1.PluginCapabilitiesClient)(nil), client.capabilitiesClient)
 	assert.Implements(t, (*storage_v1.DependenciesReaderPluginClient)(nil), client.depsReaderClient)
 	assert.Implements(t, (*storage_v1.StreamingSpanWriterPluginClient)(nil), client.streamWriterClient)
-	assert.Implements(t, (*storage_v1.SamplingStorePluginClient)(nil), client.samplingStoreClient)
+	assert.Implements(t, (*storage_v1.SamplingStorePluginClient)(nil), client.streamWriterClient)
 }
 
 func TestGRPCClientGetServices(t *testing.T) {