github.com/yandex/pandora@v0.5.32/components/providers/http/provider/provider_test.go (about)

     1  package provider
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  	"github.com/stretchr/testify/mock"
    11  	"github.com/stretchr/testify/require"
    12  	"github.com/yandex/pandora/components/providers/http/config"
    13  	"github.com/yandex/pandora/components/providers/http/decoders"
    14  	"github.com/yandex/pandora/components/providers/http/decoders/ammo"
    15  )
    16  
    17  func TestProvider_runPreloaded(t *testing.T) {
    18  	var mustNewAmmo = func(t *testing.T, method string, url string, body []byte, header http.Header, tag string) *ammo.Ammo {
    19  		a := ammo.Ammo{}
    20  		err := a.Setup(method, url, body, header, tag)
    21  		require.NoError(t, err)
    22  		return &a
    23  	}
    24  
    25  	tests := []struct {
    26  		name           string
    27  		contextTimeout time.Duration
    28  		cfg            config.Config
    29  		wantAmmos      int
    30  		wantErr        string
    31  	}{
    32  		{
    33  			name:           "ammo limit",
    34  			contextTimeout: time.Second,
    35  			cfg: config.Config{
    36  				Passes: 0,
    37  				Limit:  5,
    38  			},
    39  			wantAmmos: 5,
    40  			wantErr:   decoders.ErrAmmoLimit.Error(),
    41  		},
    42  		{
    43  			name:           "pass limit",
    44  			contextTimeout: time.Second,
    45  			cfg: config.Config{
    46  				Passes: 1,
    47  				Limit:  9,
    48  			},
    49  			wantAmmos: 3,
    50  			wantErr:   decoders.ErrPassLimit.Error(),
    51  		},
    52  		{
    53  			name:           "context deadline exceeded",
    54  			contextTimeout: 15 * time.Millisecond,
    55  			cfg: config.Config{
    56  				Passes: 1,
    57  				Limit:  9,
    58  			},
    59  			wantAmmos: 2,
    60  			wantErr:   "error from context: context deadline exceeded",
    61  		},
    62  	}
    63  	for _, tt := range tests {
    64  		t.Run(tt.name, func(t *testing.T) {
    65  			provider := &Provider{
    66  				ammos: []decoders.DecodedAmmo{
    67  					mustNewAmmo(t, "GET", "get", nil, nil, "tagget"),
    68  					mustNewAmmo(t, "POST", "post", nil, nil, "tagpost"),
    69  					mustNewAmmo(t, "PUT", "put", nil, nil, "tagput"),
    70  				},
    71  				Config: tt.cfg,
    72  				Sink:   make(chan decoders.DecodedAmmo),
    73  			}
    74  
    75  			ctx, cancel := context.WithTimeout(context.Background(), tt.contextTimeout)
    76  			defer cancel()
    77  
    78  			i := 0
    79  			cl := make(chan struct{})
    80  			go func() {
    81  				defer close(cl)
    82  				for a := range provider.Sink {
    83  					req, err := a.BuildRequest()
    84  					assert.NoError(t, err)
    85  					switch i % 3 {
    86  					case 0:
    87  						require.Equal(t, "GET", req.Method)
    88  					case 1:
    89  						require.Equal(t, "POST", req.Method)
    90  					case 2:
    91  						require.Equal(t, "PUT", req.Method)
    92  					}
    93  					i++
    94  					time.Sleep(10 * time.Millisecond) // for test context deadline exceeded
    95  				}
    96  			}()
    97  
    98  			err := provider.runPreloaded(ctx)
    99  			assert.EqualError(t, err, tt.wantErr)
   100  
   101  			close(provider.Sink)
   102  			<-cl
   103  			assert.Equal(t, tt.wantAmmos, i)
   104  		})
   105  	}
   106  
   107  }
   108  
   109  func TestLoadAmmo_ChosenTags(t *testing.T) {
   110  	var mustNewAmmo = func(t *testing.T, method string, url string, body []byte, header http.Header, tag string) *ammo.Ammo {
   111  		a := ammo.Ammo{}
   112  		err := a.Setup(method, url, body, header, tag)
   113  		require.NoError(t, err)
   114  		return &a
   115  	}
   116  	ammo1 := mustNewAmmo(t, "GET", "", nil, make(http.Header), "tag1")
   117  	ammo2 := mustNewAmmo(t, "PUT", "", nil, make(http.Header), "tag2")
   118  	ammo3 := mustNewAmmo(t, "POST", "", nil, make(http.Header), "tag3")
   119  
   120  	decoder := decoders.NewMockDecoder(t)
   121  	decoder.On("LoadAmmo", mock.Anything).Return([]decoders.DecodedAmmo{ammo1, ammo2, ammo3}, nil)
   122  
   123  	provider := &Provider{
   124  		Decoder: decoder,
   125  		ammos:   nil,
   126  		Config:  config.Config{ChosenCases: []string{"tag1", "tag3"}},
   127  	}
   128  
   129  	err := provider.loadAmmo(context.Background())
   130  	assert.NoError(t, err)
   131  
   132  	expectedAmmos := []decoders.DecodedAmmo{ammo1, ammo3}
   133  
   134  	assert.Len(t, provider.ammos, len(expectedAmmos))
   135  	assert.Equal(t, provider.ammos, expectedAmmos)
   136  }