Skip to content

Commit 435f6c4

Browse files
mmelogranommelograno
committed
added FetchMany
1 parent 4fdb9d7 commit 435f6c4

File tree

6 files changed

+108
-15
lines changed

6 files changed

+108
-15
lines changed

‎storage/inmemory/mutexmap/rulebasedsegment_test.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,13 @@ func TestRuleBasedSegmentsStorage(t *testing.T) {
8989
assert.Contains(t, names, "rule1")
9090
assert.Contains(t, names, "rule2")
9191

92+
// Test FetchMany
93+
rbsFetchMany := storage.FetchMany([]string{"rule1", "rule2", "nonexistent"})
94+
assert.Len(t, rbsFetchMany, 3)
95+
assert.Equal(t, "rule1", rbsFetchMany["rule1"].Name)
96+
assert.Equal(t, "rule2", rbsFetchMany["rule2"].Name)
97+
assert.Nil(t, rbsFetchMany["nonexistent"])
98+
9299
// Test GetSegments
93100
segments := storage.Segments()
94101
assert.True(t, segments.Has("segment1"), "segment1 should be in segments")

‎storage/interfaces.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -281,6 +281,7 @@ type RuleBasedSegmentStorageProducer interface {
281281
type RuleBasedSegmentStorageConsumer interface {
282282
ChangeNumber() (int64, error)
283283
All() []dtos.RuleBasedSegmentDTO
284+
FetchMany(rbsNames []string) map[string]*dtos.RuleBasedSegmentDTO
284285
RuleBasedSegmentNames() ([]string, error)
285286
Contains(ruleBasedSegmentNames []string) bool
286287
Segments() *set.ThreadUnsafeSet

‎storage/mocks/rulebasedsegment.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,4 +81,9 @@ func (m *MockRuleBasedSegmentStorage) LargeSegments() *set.ThreadUnsafeSet {
8181
return args.Get(0).(*set.ThreadUnsafeSet)
8282
}
8383

84+
func (m *MockRuleBasedSegmentStorage) FetchMany(rbsNames []string) map[string]*dtos.RuleBasedSegmentDTO {
85+
args := m.Called(rbsNames)
86+
return args.Get(0).(map[string]*dtos.RuleBasedSegmentDTO)
87+
}
88+
8489
var _ storage.RuleBasedSegmentsStorage = (*MockRuleBasedSegmentStorage)(nil)

‎storage/redis/rulebasedsegment_test.go‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -605,6 +605,54 @@ func TestReplaceAllRuleBased(t *testing.T) {
605605
redisClient.Del(keys...)
606606
}
607607

608+
func TestRBFetchMany(t *testing.T) {
609+
t.Run("FetchMany Error", func(t *testing.T) {
610+
expectedKey := "someprefix.SPLITIO.rbsegment.someRB1"
611+
expectedKey2 := "someprefix.SPLITIO.rbsegment.someRB2"
612+
613+
mockedRedisClient := mocks.MockClient{
614+
MGetCall: func(keys []string) redis.Result {
615+
assert.ElementsMatch(t, []string{expectedKey, expectedKey2}, keys)
616+
return &mocks.MockResultOutput{
617+
MultiInterfaceCall: func() ([]interface{}, error) {
618+
return []interface{}{}, errors.New("Some Error")
619+
},
620+
}
621+
},
622+
}
623+
mockPrefixedClient, _ := redis.NewPrefixedRedisClient(&mockedRedisClient, "someprefix")
624+
rbStorage := NewRuleBasedStorage(mockPrefixedClient, logging.NewLogger(&logging.LoggerOptions{}))
625+
rbs := rbStorage.FetchMany([]string{"someRB1", "someRB2"})
626+
assert.Nil(t, rbs)
627+
})
628+
629+
t.Run("FetchMany Success", func(t *testing.T) {
630+
expectedKey := "someprefix.SPLITIO.rbsegment.someRB1"
631+
expectedKey2 := "someprefix.SPLITIO.rbsegment.someRB2"
632+
633+
mockedRedisClient := mocks.MockClient{
634+
MGetCall: func(keys []string) redis.Result {
635+
assert.ElementsMatch(t, []string{expectedKey, expectedKey2}, keys)
636+
return &mocks.MockResultOutput{
637+
MultiInterfaceCall: func() ([]interface{}, error) {
638+
return []interface{}{
639+
marshalRuleBasedSegment(createSampleRBSegment("someRB1")),
640+
marshalRuleBasedSegment(createSampleRBSegment("someRB2")),
641+
}, nil
642+
},
643+
}
644+
},
645+
}
646+
mockPrefixedClient, _ := redis.NewPrefixedRedisClient(&mockedRedisClient, "someprefix")
647+
648+
rbStorage := NewRuleBasedStorage(mockPrefixedClient, logging.NewLogger(&logging.LoggerOptions{}))
649+
rbs := rbStorage.FetchMany([]string{"someRB1", "someRB2"})
650+
assert.Equal(t, 2, len(rbs))
651+
assert.NotNil(t, rbs["someRB1"])
652+
assert.NotNil(t, rbs["someRB2"])
653+
})
654+
}
655+
608656
func marshalRuleBasedSegment(rbSegment dtos.RuleBasedSegmentDTO) string {
609657
json, _ := json.Marshal(rbSegment)
610658
return string(json)

‎storage/redis/rulebasedsegments.go‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -368,4 +368,36 @@ func (r *RuleBasedSegmentStorage) executePipeline(pipeline redis.Pipeline, toAdd
368368
return failedToAdd, failedToRemove
369369
}
370370

371+
func (r *RuleBasedSegmentStorage) FetchMany(names []string) map[string]*dtos.RuleBasedSegmentDTO {
372+
if len(names) == 0 {
373+
return nil
374+
}
375+
376+
keysToFetch := make([]string, 0, len(names))
377+
for _, name := range names {
378+
keysToFetch = append(keysToFetch, strings.Replace(KeyRuleBasedSegment, "{rbsegment}", name, 1))
379+
}
380+
rawRBS, err := r.client.MGet(keysToFetch)
381+
if err != nil {
382+
r.logger.Error(fmt.Sprintf("Could not fetch rule-based segments from redis: %s", err.Error()))
383+
return nil
384+
}
385+
386+
rbs := make(map[string]*dtos.RuleBasedSegmentDTO)
387+
for idx, rb := range names {
388+
var rbSegment *dtos.RuleBasedSegmentDTO
389+
rawRBSegment, ok := rawRBS[idx].(string)
390+
if ok {
391+
err = json.Unmarshal([]byte(rawRBSegment), &rbSegment)
392+
if err != nil {
393+
r.logger.Error("Could not parse rule-based segment \"%s\" fetched from redis", rb)
394+
return nil
395+
}
396+
}
397+
rbs[rb] = rbSegment
398+
}
399+
400+
return rbs
401+
}
402+
371403
var _ storage.RuleBasedSegmentsStorage = (*RuleBasedSegmentStorage)(nil)

‎storage/redis/splits_test.go‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -562,20 +562,20 @@ func TestUpdateRedis(t *testing.T) {
562562
if len(splits) != 3 {
563563
t.Error("Unexpected amount of splits")
564564
}
565-
set1, err := redisClient.SMembers("SPLITIO.flagSet.set1")
565+
set1, _ := redisClient.SMembers("SPLITIO.flagSet.set1")
566566
if len(set1) != 2 {
567567
t.Error("set size should be 2")
568568
}
569569
if !slices.Contains(set1, "split1") || !slices.Contains(set1, "split2") {
570570
t.Error("Split missing in set")
571571
}
572-
tt, err := redisClient.Get("SPLITIO.trafficType.user")
573-
ttCount, _ := strconv.ParseFloat(tt, 10)
572+
tt, _ := redisClient.Get("SPLITIO.trafficType.user")
573+
ttCount, _ := strconv.ParseFloat(tt, 64)
574574
if ttCount != 3 {
575575
t.Error("Split should exist")
576576
}
577-
till, err := redisClient.Get("SPLITIO.splits.till")
578-
tillInt, _ := strconv.ParseFloat(till, 10)
577+
till, _ := redisClient.Get("SPLITIO.splits.till")
578+
tillInt, _ := strconv.ParseFloat(till, 64)
579579
if tillInt != 1 {
580580
t.Error("ChangeNumber should be 1")
581581
}
@@ -587,29 +587,29 @@ func TestUpdateRedis(t *testing.T) {
587587
if len(splits) != 3 {
588588
t.Error("Unexpected size")
589589
}
590-
set1, err = redisClient.SMembers("SPLITIO.flagSet.set1")
590+
set1, _ = redisClient.SMembers("SPLITIO.flagSet.set1")
591591
if len(set1) != 0 {
592592
t.Error("set size should be 0")
593593
}
594-
set3, err := redisClient.SMembers("SPLITIO.flagSet.set3")
594+
set3, _ := redisClient.SMembers("SPLITIO.flagSet.set3")
595595
if len(set3) != 3 {
596596
t.Error("set size should be 3")
597597
}
598598
if !slices.Contains(set3, "split3") || !slices.Contains(set3, "split4") || !slices.Contains(set3, "split5") {
599599
t.Error("Split missing in set")
600600
}
601-
tt, err = redisClient.Get("SPLITIO.trafficType.user")
602-
ttCount, _ = strconv.ParseFloat(tt, 10)
601+
tt, _ = redisClient.Get("SPLITIO.trafficType.user")
602+
ttCount, _ = strconv.ParseFloat(tt, 64)
603603
if ttCount != 3 {
604604
t.Error("Unexpected trafficType occurrences")
605605
}
606606

607-
split1, err := redisClient.Get("SPLITIO.split.split1")
607+
split1, _ := redisClient.Get("SPLITIO.split.split1")
608608
if split1 != "" {
609609
t.Error("Split should not exist")
610610
}
611-
till, err = redisClient.Get("SPLITIO.splits.till")
612-
tillInt, _ = strconv.ParseFloat(till, 10)
611+
till, _ = redisClient.Get("SPLITIO.splits.till")
612+
tillInt, _ = strconv.ParseFloat(till, 64)
613613
if tillInt != 2 {
614614
t.Error("ChangeNumber should be 2")
615615
}
@@ -649,15 +649,15 @@ func TestUpdateWithFlagSetFiltersRedis(t *testing.T) {
649649
if len(splits) != 3 {
650650
t.Error("Unexpected amount of splits")
651651
}
652-
set1, err := redisClient.SMembers("SPLITIO.flagSet.set1")
652+
set1, _ := redisClient.SMembers("SPLITIO.flagSet.set1")
653653
if len(set1) != 2 {
654654
t.Error("set size should be 2")
655655
}
656-
set2, err := redisClient.SMembers("SPLITIO.flagSet.set2")
656+
set2, _ := redisClient.SMembers("SPLITIO.flagSet.set2")
657657
if len(set2) != 1 {
658658
t.Error("set size should be 1")
659659
}
660-
set3, err := redisClient.SMembers("SPLITIO.flagSet.set3")
660+
set3, _ := redisClient.SMembers("SPLITIO.flagSet.set3")
661661
if len(set3) != 0 {
662662
t.Error("set size should be 0")
663663
}

0 commit comments

Comments
 (0)