github.com/thanos-io/thanos@v0.32.5/pkg/api/query/grpc_test.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package v1
     5  
     6  import (
     7  	"context"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/efficientgo/core/errors"
    12  	"github.com/efficientgo/core/testutil"
    13  	"github.com/go-kit/log"
    14  	"github.com/prometheus/client_golang/prometheus"
    15  	"github.com/prometheus/prometheus/promql"
    16  	"github.com/prometheus/prometheus/storage"
    17  	v1 "github.com/prometheus/prometheus/web/api/v1"
    18  
    19  	"github.com/thanos-io/thanos/pkg/api/query/querypb"
    20  	"github.com/thanos-io/thanos/pkg/component"
    21  	"github.com/thanos-io/thanos/pkg/query"
    22  	"github.com/thanos-io/thanos/pkg/store"
    23  )
    24  
    25  func TestGRPCQueryAPIErrorHandling(t *testing.T) {
    26  	logger := log.NewNopLogger()
    27  	reg := prometheus.NewRegistry()
    28  	proxy := store.NewProxyStore(logger, reg, func() []store.Client { return nil }, component.Store, nil, 1*time.Minute, store.LazyRetrieval)
    29  	queryableCreator := query.NewQueryableCreator(logger, reg, proxy, 1, 1*time.Minute)
    30  	lookbackDeltaFunc := func(i int64) time.Duration { return 5 * time.Minute }
    31  	tests := []struct {
    32  		name   string
    33  		engine *engineStub
    34  	}{
    35  		{
    36  			name: "error response",
    37  			engine: &engineStub{
    38  				err: errors.New("error stub"),
    39  			},
    40  		},
    41  		{
    42  			name: "error response",
    43  			engine: &engineStub{
    44  				warns: []error{errors.New("warn stub")},
    45  			},
    46  		},
    47  	}
    48  
    49  	for _, test := range tests {
    50  		engineFactory := &QueryEngineFactory{
    51  			prometheusEngine: test.engine,
    52  		}
    53  		api := NewGRPCAPI(time.Now, nil, queryableCreator, engineFactory, querypb.EngineType_prometheus, lookbackDeltaFunc, 0)
    54  		t.Run("range_query", func(t *testing.T) {
    55  			rangeRequest := &querypb.QueryRangeRequest{
    56  				Query:            "metric",
    57  				StartTimeSeconds: 0,
    58  				IntervalSeconds:  10,
    59  				EndTimeSeconds:   300,
    60  			}
    61  			srv := newQueryRangeServer(context.Background())
    62  			err := api.QueryRange(rangeRequest, srv)
    63  
    64  			if test.engine.err != nil {
    65  				testutil.NotOk(t, err)
    66  				return
    67  			}
    68  			if len(test.engine.warns) > 0 {
    69  				testutil.Ok(t, err)
    70  				for i, resp := range srv.responses {
    71  					testutil.Equals(t, test.engine.warns[i].Error(), resp.GetWarnings())
    72  				}
    73  			}
    74  		})
    75  
    76  		t.Run("instant_query", func(t *testing.T) {
    77  			instantRequest := &querypb.QueryRequest{
    78  				Query:          "metric",
    79  				TimeoutSeconds: 60,
    80  			}
    81  			srv := newQueryServer(context.Background())
    82  			err := api.Query(instantRequest, srv)
    83  			if test.engine.err != nil {
    84  				testutil.NotOk(t, err)
    85  				return
    86  			}
    87  			if len(test.engine.warns) > 0 {
    88  				testutil.Ok(t, err)
    89  				for i, resp := range srv.responses {
    90  					testutil.Equals(t, test.engine.warns[i].Error(), resp.GetWarnings())
    91  				}
    92  			}
    93  		})
    94  	}
    95  }
    96  
    97  type engineStub struct {
    98  	v1.QueryEngine
    99  	err   error
   100  	warns []error
   101  }
   102  
   103  func (e engineStub) NewInstantQuery(_ context.Context, q storage.Queryable, opts promql.QueryOpts, qs string, ts time.Time) (promql.Query, error) {
   104  	return &queryStub{err: e.err, warns: e.warns}, nil
   105  }
   106  
   107  func (e engineStub) NewRangeQuery(_ context.Context, q storage.Queryable, opts promql.QueryOpts, qs string, start, end time.Time, interval time.Duration) (promql.Query, error) {
   108  	return &queryStub{err: e.err, warns: e.warns}, nil
   109  }
   110  
   111  type queryStub struct {
   112  	promql.Query
   113  	err   error
   114  	warns []error
   115  }
   116  
   117  func (q queryStub) Close() {}
   118  
   119  func (q queryStub) Exec(context.Context) *promql.Result {
   120  	return &promql.Result{Err: q.err, Warnings: q.warns}
   121  }
   122  
   123  type queryServer struct {
   124  	querypb.Query_QueryServer
   125  
   126  	ctx       context.Context
   127  	responses []querypb.QueryResponse
   128  }
   129  
   130  func newQueryServer(ctx context.Context) *queryServer {
   131  	return &queryServer{ctx: ctx}
   132  }
   133  
   134  func (q *queryServer) Send(r *querypb.QueryResponse) error {
   135  	q.responses = append(q.responses, *r)
   136  	return nil
   137  }
   138  
   139  func (q *queryServer) Context() context.Context {
   140  	return q.ctx
   141  }
   142  
   143  type queryRangeServer struct {
   144  	querypb.Query_QueryRangeServer
   145  
   146  	ctx       context.Context
   147  	responses []querypb.QueryRangeResponse
   148  }
   149  
   150  func newQueryRangeServer(ctx context.Context) *queryRangeServer {
   151  	return &queryRangeServer{ctx: ctx}
   152  }
   153  
   154  func (q *queryRangeServer) Send(r *querypb.QueryRangeResponse) error {
   155  	q.responses = append(q.responses, *r)
   156  	return nil
   157  }
   158  
   159  func (q *queryRangeServer) Context() context.Context {
   160  	return q.ctx
   161  }