add suggested changes

Signed-off-by: slayer321 <sachin.maurya7666@gmail.com>
jaegertracing · Oct 16, 2023 · 04f67c4 · 04f67c4
1 parent 6a3ee0d
commit 04f67c4
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 32 deletions.
diff --git a/plugin/storage/badger/factory.go b/plugin/storage/badger/factory.go
@@ -29,7 +29,7 @@ import (
 	"github.com/jaegertracing/jaeger/pkg/metrics"
 	"github.com/jaegertracing/jaeger/plugin"
 	depStore "github.com/jaegertracing/jaeger/plugin/storage/badger/dependencystore"
-	samplingStore "github.com/jaegertracing/jaeger/plugin/storage/badger/samplingstore"
+	badgerSampling "github.com/jaegertracing/jaeger/plugin/storage/badger/samplingstore"
 	badgerStore "github.com/jaegertracing/jaeger/plugin/storage/badger/spanstore"
 	"github.com/jaegertracing/jaeger/storage/dependencystore"
 	"github.com/jaegertracing/jaeger/storage/samplingstore"
@@ -174,7 +174,7 @@ func (f *Factory) CreateDependencyReader() (dependencystore.Reader, error) {
 
 // CreateSamplingStore implements storage.SamplingStoreFactory
 func (f *Factory) CreateSamplingStore(maxBuckets int) (samplingstore.Store, error) {
-	return samplingStore.NewSamplingStore(f.store), nil
+	return badgerSampling.NewSamplingStore(f.store), nil
 }
 
 // Close Implements io.Closer and closes the underlying storage

diff --git a/plugin/storage/badger/samplingstore/storage_test.go b/plugin/storage/badger/samplingstore/storage_test.go
@@ -15,6 +15,7 @@
 package samplingstore
 
 import (
+	"encoding/json"
 	"testing"
 	"time"
 
@@ -24,67 +25,106 @@ import (
 	samplemodel "github.com/jaegertracing/jaeger/cmd/collector/app/sampling/model"
 )
 
-type samplingStoreTest struct {
-	store *SamplingStore
-}
-
-func NewtestSamplingStore(db *badger.DB) *samplingStoreTest {
-	return &samplingStoreTest{
-		store: NewSamplingStore(db),
-	}
+func newTestSamplingStore(db *badger.DB) *SamplingStore {
+	return NewSamplingStore(db)
 }
 
 func TestInsertThroughput(t *testing.T) {
-	runWithBadger(t, func(s *samplingStoreTest, t *testing.T) {
+	runWithBadger(t, func(t *testing.T, store *SamplingStore) {
 		throughputs := []*samplemodel.Throughput{
 			{Service: "my-svc", Operation: "op"},
 			{Service: "our-svc", Operation: "op2"},
 		}
-		err := s.store.InsertThroughput(throughputs)
+		err := store.InsertThroughput(throughputs)
 		assert.NoError(t, err)
 	})
 }
 
 func TestGetThroughput(t *testing.T) {
-	runWithBadger(t, func(s *samplingStoreTest, t *testing.T) {
+	runWithBadger(t, func(t *testing.T, store *SamplingStore) {
 		start := time.Now()
-
-		expected := 2
-		throughputs := []*samplemodel.Throughput{
+		expected := []*samplemodel.Throughput{
 			{Service: "my-svc", Operation: "op"},
 			{Service: "our-svc", Operation: "op2"},
 		}
-		err := s.store.InsertThroughput(throughputs)
+		err := store.InsertThroughput(expected)
 		assert.NoError(t, err)
 
-		actual, err := s.store.GetThroughput(start, start.Add(time.Second*time.Duration(10)))
+		actual, err := store.GetThroughput(start, start.Add(time.Second*time.Duration(10)))
 		assert.NoError(t, err)
-		assert.Equal(t, expected, len(actual))
+		assert.Equal(t, expected, actual)
 	})
 }
 
 func TestInsertProbabilitiesAndQPS(t *testing.T) {
-	runWithBadger(t, func(s *samplingStoreTest, t *testing.T) {
-		err := s.store.InsertProbabilitiesAndQPS("dell11eg843d", samplemodel.ServiceOperationProbabilities{"new-srv": {"op": 0.1}}, samplemodel.ServiceOperationQPS{"new-srv": {"op": 4}})
+	runWithBadger(t, func(t *testing.T, store *SamplingStore) {
+		err := store.InsertProbabilitiesAndQPS(
+			"dell11eg843d",
+			samplemodel.ServiceOperationProbabilities{"new-srv": {"op": 0.1}},
+			samplemodel.ServiceOperationQPS{"new-srv": {"op": 4}},
+		)
 		assert.NoError(t, err)
 	})
 }
 
 func TestGetLatestProbabilities(t *testing.T) {
-	runWithBadger(t, func(s *samplingStoreTest, t *testing.T) {
-		err := s.store.InsertProbabilitiesAndQPS("dell11eg843d", samplemodel.ServiceOperationProbabilities{"new-srv": {"op": 0.1}}, samplemodel.ServiceOperationQPS{"new-srv": {"op": 4}})
+	runWithBadger(t, func(t *testing.T, store *SamplingStore) {
+		err := store.InsertProbabilitiesAndQPS(
+			"dell11eg843d",
+			samplemodel.ServiceOperationProbabilities{"new-srv": {"op": 0.1}},
+			samplemodel.ServiceOperationQPS{"new-srv": {"op": 4}},
+		)
 		assert.NoError(t, err)
-		err = s.store.InsertProbabilitiesAndQPS("newhostname", samplemodel.ServiceOperationProbabilities{"new-srv2": {"op": 0.123}}, samplemodel.ServiceOperationQPS{"new-srv2": {"op": 1}})
+		err = store.InsertProbabilitiesAndQPS(
+			"newhostname",
+			samplemodel.ServiceOperationProbabilities{"new-srv2": {"op": 0.123}},
+			samplemodel.ServiceOperationQPS{"new-srv2": {"op": 1}},
+		)
 		assert.NoError(t, err)
 
 		expected := samplemodel.ServiceOperationProbabilities{"new-srv2": {"op": 0.123}}
-		actual, err := s.store.GetLatestProbabilities()
+		actual, err := store.GetLatestProbabilities()
 		assert.NoError(t, err)
 		assert.Equal(t, expected, actual)
 	})
 }
 
-func runWithBadger(t *testing.T, test func(s *samplingStoreTest, t *testing.T)) {
+func TestDecodeProbabilitiesValue(t *testing.T) {
+	expected := ProbabilitiesAndQPS{
+		Hostname:      "dell11eg843d",
+		Probabilities: samplemodel.ServiceOperationProbabilities{"new-srv": {"op": 0.1}},
+		QPS:           samplemodel.ServiceOperationQPS{"new-srv": {"op": 4}},
+	}
+
+	marshalBytes, err := json.Marshal(expected)
+	assert.NoError(t, err)
+	actual, err := decodeProbabilitiesValue(marshalBytes)
+	assert.NoError(t, err)
+	assert.Equal(t, expected, actual)
+}
+
+func TestDecodeThroughtputValue(t *testing.T) {
+	expected := []*samplemodel.Throughput{
+		{Service: "my-svc", Operation: "op"},
+		{Service: "our-svc", Operation: "op2"},
+	}
+
+	marshalBytes, err := json.Marshal(expected)
+	assert.NoError(t, err)
+	acrual, err := decodeThroughtputValue(marshalBytes)
+	assert.NoError(t, err)
+	assert.Equal(t, expected, acrual)
+}
+
+func TestInitalStartTime(t *testing.T) {
+	expected := "2023-10-10 20:17:19.993838 +0530 IST"
+	timeBytes := []byte{0, 6, 7, 93, 200, 166, 213, 238}
+	actual, err := initalStartTime(timeBytes)
+	assert.NoError(t, err)
+	assert.Equal(t, expected, actual.String())
+}
+
+func runWithBadger(t *testing.T, test func(t *testing.T, store *SamplingStore)) {
 	opts := badger.DefaultOptions("")
 
 	opts.SyncWrites = false
@@ -94,11 +134,11 @@ func runWithBadger(t *testing.T, test func(s *samplingStoreTest, t *testing.T))
 
 	store, err := badger.Open(opts)
 	defer func() {
-		store.Close()
+		assert.NoError(t, store.Close())
 	}()
-	ss := NewtestSamplingStore(store)
+	ss := newTestSamplingStore(store)
 
 	assert.NoError(t, err)
 
-	test(ss, t)
+	test(t, ss)
 }
diff --git a/plugin/storage/integration/integration.go b/plugin/storage/integration/integration.go
@@ -433,15 +433,18 @@ func (s *StorageIntegration) testGetThroughput(t *testing.T) {
 
 	s.insertThroughput(t)
 
-	expected := 2
+	expected := []*samplemodel.Throughput{
+		{Service: "my-svc", Operation: "op"},
+		{Service: "our-svc", Operation: "op2"},
+	}
 	var actual []*samplemodel.Throughput
 	_ = s.waitForCondition(t, func(t *testing.T) bool {
 		var err error
 		actual, err = s.SamplingStore.GetThroughput(start, start.Add(time.Second*time.Duration(10)))
 		require.NoError(t, err)
-		return assert.ObjectsAreEqualValues(expected, len(actual))
+		return assert.ObjectsAreEqualValues(expected, actual)
 	})
-	assert.Len(t, actual, expected)
+	assert.Equal(t, expected, actual)
 }
 
 func (s *StorageIntegration) testGetLatestProbability(t *testing.T) {