github.com/grafana/synthetic-monitoring-agent@v0.64.2/internal/pkg/prom/prom_test.go (about)

     1  package prom_test
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"errors"
     7  	"io"
     8  	"testing"
     9  
    10  	"github.com/grafana/synthetic-monitoring-agent/internal/pkg/prom"
    11  )
    12  
    13  type FailNTimesPrometheusClient struct {
    14  	CountedRetries float64
    15  	FailuresLeft   int64
    16  }
    17  
    18  func (f *FailNTimesPrometheusClient) StoreBytes(ctx context.Context, req []byte) error {
    19  	return f.StoreStream(ctx, bytes.NewReader(req))
    20  }
    21  
    22  func (f *FailNTimesPrometheusClient) StoreStream(ctx context.Context, req io.Reader) error {
    23  	if f.FailuresLeft == 0 {
    24  		return nil
    25  	}
    26  
    27  	f.FailuresLeft--
    28  
    29  	return prom.NewRecoverableError(errors.New("client failed"))
    30  }
    31  
    32  func (f *FailNTimesPrometheusClient) CountRetries(retries float64) {
    33  	f.CountedRetries = retries
    34  }
    35  
    36  type FakePrometheusClient struct {
    37  	storeFunc        func(ctx context.Context, req []byte) error
    38  	countRetriesFunc func(float642 float64)
    39  }
    40  
    41  func (f *FakePrometheusClient) Store(ctx context.Context, req []byte) error {
    42  	return f.storeFunc(ctx, req)
    43  }
    44  
    45  func (f *FakePrometheusClient) CountRetries(retries float64) {
    46  	f.countRetriesFunc(retries)
    47  }
    48  
    49  func TestSendBytesWithBackoffRetriesCounter(t *testing.T) {
    50  	type args struct {
    51  		ctx         context.Context
    52  		timesToFail int64
    53  		req         []byte
    54  	}
    55  
    56  	tests := []struct {
    57  		name                 string
    58  		args                 args
    59  		wantErr              bool
    60  		expectedRetriesCount float64
    61  		isSlow               bool
    62  	}{
    63  		{
    64  			name: "should count 0 retries when successful at first",
    65  			args: args{
    66  				ctx:         context.Background(),
    67  				timesToFail: 0,
    68  				req:         []byte{},
    69  			},
    70  			expectedRetriesCount: 0,
    71  		},
    72  		{
    73  			name: "should count 2 retries when failed 2 times",
    74  			args: args{
    75  				ctx:         context.Background(),
    76  				timesToFail: 2,
    77  				req:         []byte{},
    78  			},
    79  			expectedRetriesCount: 2,
    80  		},
    81  		{
    82  			name: "should count 10 retries when retries exceeded",
    83  			args: args{
    84  				ctx:         context.Background(),
    85  				timesToFail: 10,
    86  				req:         []byte{},
    87  			},
    88  			wantErr:              true,
    89  			expectedRetriesCount: 10,
    90  			isSlow:               true,
    91  		},
    92  	}
    93  	for _, tt := range tests {
    94  		client := &FailNTimesPrometheusClient{
    95  			CountedRetries: -1,
    96  			FailuresLeft:   tt.args.timesToFail,
    97  		}
    98  		t.Run(tt.name, func(t *testing.T) {
    99  			if testing.Short() && tt.isSlow {
   100  				t.Skip("skipping in short mode")
   101  			}
   102  
   103  			if err := prom.SendBytesWithBackoff(tt.args.ctx, client, tt.args.req); (err != nil) != tt.wantErr {
   104  				t.Errorf("SendBytesWithBackoff() error = %v, wantErr %v", err, tt.wantErr)
   105  				return
   106  			}
   107  
   108  			if client.CountedRetries != tt.expectedRetriesCount {
   109  				t.Errorf("Not expected retries count registered. Expected: %f, Got: %f", tt.expectedRetriesCount, client.CountedRetries)
   110  			}
   111  		})
   112  	}
   113  }