From 035603954a91113f258a0b7478a83c12fa20c1aa Mon Sep 17 00:00:00 2001
From: JaeguKim <rlaworn1993@gmail.com>
Date: Fri, 12 Apr 2024 19:52:31 +0900
Subject: [PATCH] Create sampling store in initializeESFactory

Signed-off-by: JaeguKim <rlaworn1993@gmail.com>
---
 .../storage/integration/elasticsearch_test.go | 43 ++-----------------
 1 file changed, 3 insertions(+), 40 deletions(-)

diff --git a/plugin/storage/integration/elasticsearch_test.go b/plugin/storage/integration/elasticsearch_test.go
index 2579d213bb8..df5992121a6 100644
--- a/plugin/storage/integration/elasticsearch_test.go
+++ b/plugin/storage/integration/elasticsearch_test.go
@@ -33,13 +33,9 @@ import (
 	"go.uber.org/zap/zaptest"
 
 	"github.com/jaegertracing/jaeger/pkg/config"
-	estemplate "github.com/jaegertracing/jaeger/pkg/es"
-	eswrapper "github.com/jaegertracing/jaeger/pkg/es/wrapper"
 	"github.com/jaegertracing/jaeger/pkg/metrics"
 	"github.com/jaegertracing/jaeger/pkg/testutils"
 	"github.com/jaegertracing/jaeger/plugin/storage/es"
-	"github.com/jaegertracing/jaeger/plugin/storage/es/mappings"
-	"github.com/jaegertracing/jaeger/plugin/storage/es/samplingstore"
 	"github.com/jaegertracing/jaeger/storage/dependencystore"
 )
 
@@ -102,7 +98,6 @@ func (s *ESStorageIntegration) initializeES(t *testing.T, allTagsAsFields bool)
 	require.NoError(t, err)
 
 	s.initSpanstore(t, allTagsAsFields)
-	s.initSamplingStore(t)
 
 	s.CleanUp = func(t *testing.T) {
 		s.esCleanUp(t, allTagsAsFields)
@@ -120,41 +115,6 @@ func (s *ESStorageIntegration) esCleanUp(t *testing.T, allTagsAsFields bool) {
 	s.initSpanstore(t, allTagsAsFields)
 }
 
-func (s *ESStorageIntegration) initSamplingStore(t *testing.T) {
-	client := s.getEsClient(t)
-	mappingBuilder := mappings.MappingBuilder{
-		TemplateBuilder: estemplate.TextTemplateBuilder{},
-		Shards:          5,
-		Replicas:        1,
-		EsVersion:       client.GetVersion(),
-		IndexPrefix:     indexPrefix,
-		UseILM:          false,
-	}
-	clientFn := func() estemplate.Client { return client }
-	samplingstore := samplingstore.NewSamplingStore(
-		samplingstore.SamplingStoreParams{
-			Client:          clientFn,
-			Logger:          s.logger,
-			IndexPrefix:     indexPrefix,
-			IndexDateLayout: indexDateLayout,
-			MaxDocCount:     defaultMaxDocCount,
-		})
-	sampleMapping, err := mappingBuilder.GetSamplingMappings()
-	require.NoError(t, err)
-	err = samplingstore.CreateTemplates(sampleMapping)
-	require.NoError(t, err)
-	s.SamplingStore = samplingstore
-}
-
-func (s *ESStorageIntegration) getEsClient(t *testing.T) eswrapper.ClientWrapper {
-	bp, err := s.client.BulkProcessor().BulkActions(1).FlushInterval(time.Nanosecond).Do(context.Background())
-	require.NoError(t, err)
-	s.bulkProcessor = bp
-	esVersion, err := s.getVersion()
-	require.NoError(t, err)
-	return eswrapper.WrapESClient(s.client, bp, esVersion, s.v8Client)
-}
-
 func (s *ESStorageIntegration) initializeESFactory(t *testing.T, allTagsAsFields bool) *es.Factory {
 	s.logger = zaptest.NewLogger(t)
 	f := es.NewFactory()
@@ -193,6 +153,9 @@ func (s *ESStorageIntegration) initSpanstore(t *testing.T, allTagsAsFields bool)
 	s.DependencyReader, err = f.CreateDependencyReader()
 	require.NoError(t, err)
 	s.DependencyWriter = s.DependencyReader.(dependencystore.Writer)
+
+	s.SamplingStore, err = f.CreateSamplingStore(1)
+	require.NoError(t, err)
 }
 
 func (s *ESStorageIntegration) esRefresh(t *testing.T) {