Fix tests

Signed-off-by: Kemal Akkoyun <[email protected]>
thanos-io · Jun 5, 2020 · 7e899f9 · 7e899f9
1 parent f5948ba
commit 7e899f9
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/pkg/query/iter.go b/pkg/query/iter.go
@@ -676,7 +676,7 @@ type asyncSeriesSet struct {
 	result  storage.SeriesSet
 }
 
-func newAsyncSeriesSet(ctx context.Context, gate *gate.Gate, f func() (storage.SeriesSet, storage.Warnings, error)) storage.SeriesSet {
+func newAsyncSeriesSet(ctx context.Context, gate *gate.Gate, f func(ctx context.Context) (storage.SeriesSet, storage.Warnings, error)) storage.SeriesSet {
 	promise := make(chan storage.SeriesSet, 1)
 	go func() {
 		defer close(promise)
@@ -686,7 +686,7 @@ func newAsyncSeriesSet(ctx context.Context, gate *gate.Gate, f func() (storage.S
 		}
 		defer gate.Done()
 
-		set, _, err := f()
+		set, _, err := f(ctx)
 		// TODO(kakkoyun): Handle warnings after Prometheus changes.
 		if err != nil {
 			promise <- storage.ErrSeriesSet(err)

diff --git a/pkg/query/querier.go b/pkg/query/querier.go
@@ -176,12 +176,12 @@ func (q *querier) Select(_ bool, hints *storage.SelectHints, ms ...*labels.Match
 	}
 
 	// TODO(kakkoyun): Introduce a flag for maxConcurrentSelects?
-	return newAsyncSeriesSet(q.ctx, gate.NewGate(4, q.reg), func() (storage.SeriesSet, storage.Warnings, error) {
+	return newAsyncSeriesSet(q.ctx, gate.NewGate(4, q.reg), func(ctx context.Context) (storage.SeriesSet, storage.Warnings, error) {
 		matchers := make([]string, len(ms))
 		for i, m := range ms {
 			matchers[i] = m.String()
 		}
-		span, ctx := tracing.StartSpan(q.ctx, "querier_select", opentracing.Tags{
+		span, ctx := tracing.StartSpan(ctx, "querier_select", opentracing.Tags{
 			"minTime":  hints.Start,
 			"maxTime":  hints.End,
 			"matchers": "{" + strings.Join(matchers, ",") + "}",

diff --git a/pkg/query/querier_test.go b/pkg/query/querier_test.go
@@ -506,12 +506,16 @@ func TestQuerier_Select(t *testing.T) {
 				{dedup: false, expected: tcase.expected},
 				{dedup: true, expected: []series{tcase.expectedAfterDedup}},
 			} {
-				q := newQuerier(context.Background(), nil, nil, tcase.mint, tcase.maxt, tcase.replicaLabels, tcase.storeAPI, sc.dedup, 0, true, false)
+				timeout := 10 * time.Second
+				ctx, cancel := context.WithTimeout(context.Background(), timeout)
+				t.Cleanup(cancel)
+
+				q := newQuerier(ctx, nil, nil, tcase.mint, tcase.maxt, tcase.replicaLabels, tcase.storeAPI, sc.dedup, 0, true, false)
 				defer testutil.Ok(t, q.Close())
 
 				t.Run(fmt.Sprintf("dedup=%v", sc.dedup), func(t *testing.T) {
 					t.Run("querier.Select", func(t *testing.T) {
-						defer leaktest.CheckTimeout(t, 10*time.Second)()
+						defer leaktest.CheckTimeout(t, timeout)()
 						defer func() { testutil.Ok(t, q.Close()) }()
 
 						res := q.Select(false, tcase.hints, tcase.matchers...)