test: fix tests

grafana · Sep 18, 2024 · 0a0559a · 0a0559a
1 parent b2a7688
commit 0a0559a
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
diff --git a/pkg/querier/queryrange/roundtrip_test.go b/pkg/querier/queryrange/roundtrip_test.go
@@ -1607,10 +1607,21 @@ func indexStatsResult(v logproto.IndexStatsResponse) (*int, base.Handler) {
 func seriesVolumeResult(v logproto.VolumeResponse) (*int, base.Handler) {
 	count := 0
 	var lock sync.Mutex
-	return &count, base.HandlerFunc(func(_ context.Context, _ base.Request) (base.Response, error) {
+	return &count, base.HandlerFunc(func(_ context.Context, req base.Request) (base.Response, error) {
 		lock.Lock()
 		defer lock.Unlock()
-		count++
+
+		switch req.(type) {
+		case *logproto.VolumeRequest:
+			count++
+			// Split by label middleware makes a labels call for __stream_shard__. Pretend there's never
+			// and stream shards to keep the tests simpler.
+		case *LabelRequest:
+			return &LokiLabelNamesResponse{
+				Status: "success",
+				Data:   []string{},
+			}, nil
+		}
 		return &VolumeResponse{Response: &v}, nil
 	})
 }

diff --git a/pkg/querier/queryrange/split_by_shard_label_test.go b/pkg/querier/queryrange/split_by_shard_label_test.go
@@ -160,7 +160,7 @@ func Test_SplitByShardLabel(t *testing.T) {
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			lims := limits{
-				Limits:        nil,
+				Limits:        fakeLimits{},
 				splitDuration: &tc.splitDuration,
 			}
 
@@ -228,7 +228,7 @@ func Test_SplitByShardLabel(t *testing.T) {
 		func(t *testing.T) {
 			splitDuration := 30 * time.Minute
 			lims := limits{
-				Limits:        nil,
+				Limits:        fakeLimits{},
 				splitDuration: &splitDuration,
 			}