github.com/thanos-io/thanos@v0.32.5/pkg/store/storepb/testutil/store_series_client.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package storetestutil
     5  
     6  import (
     7  	"context"
     8  	"io"
     9  	"time"
    10  
    11  	"github.com/thanos-io/thanos/pkg/store/storepb"
    12  )
    13  
    14  // StoreSeriesClient is test gRPC storeAPI series client.
    15  type StoreSeriesClient struct {
    16  	// This field just exist to pseudo-implement the unused methods of the interface.
    17  	storepb.Store_SeriesClient
    18  	Ctx             context.Context
    19  	i               int
    20  	RespSet         []*storepb.SeriesResponse
    21  	RespDur         time.Duration
    22  	SlowSeriesIndex int
    23  
    24  	InjectedError      error
    25  	InjectedErrorIndex int
    26  }
    27  
    28  func (c *StoreSeriesClient) Recv() (*storepb.SeriesResponse, error) {
    29  	if c.RespDur != 0 && (c.SlowSeriesIndex == c.i || c.SlowSeriesIndex == 0) {
    30  		select {
    31  		case <-time.After(c.RespDur):
    32  		case <-c.Ctx.Done():
    33  			return nil, c.Ctx.Err()
    34  		}
    35  	}
    36  	if c.InjectedError != nil && (c.InjectedErrorIndex == c.i || c.InjectedErrorIndex == 0) {
    37  		return nil, c.InjectedError
    38  	}
    39  
    40  	if c.i >= len(c.RespSet) {
    41  		return nil, io.EOF
    42  	}
    43  	s := c.RespSet[c.i]
    44  	c.i++
    45  
    46  	return s, nil
    47  }
    48  
    49  func (c *StoreSeriesClient) Context() context.Context {
    50  	return c.Ctx
    51  }